From 699205e0c5e12403c14c674cb47fcc5421c5d07e Mon Sep 17 00:00:00 2001 From: hallerite Date: Fri, 14 Aug 2026 15:28:01 +0200 Subject: [PATCH] fix(client): preserve session update ordering --- src/acp/client/connection.py | 78 ++++++++++++++++-- tests/test_rpc.py | 149 +++++++++++++++++++++++++++++++++++ 2 files changed, 219 insertions(+), 8 deletions(-) diff --git a/src/acp/client/connection.py b/src/acp/client/connection.py index 81f6769..000a5da 100644 --- a/src/acp/client/connection.py +++ b/src/acp/client/connection.py @@ -2,12 +2,15 @@ import asyncio from collections.abc import Callable +from contextvars import ContextVar from typing import Any, cast, final from .._transport import Transport from ..connection import Connection +from ..exceptions import RequestError from ..interfaces import Agent, Client -from ..meta import AGENT_METHODS +from ..meta import AGENT_METHODS, CLIENT_METHODS +from ..router import _resolve_handler, _warn_legacy_handler from ..schema import ( AcpMcpServer, AudioContentBlock, @@ -37,6 +40,7 @@ ResourceContentBlock, ResumeSessionRequest, ResumeSessionResponse, + SessionNotification, SetSessionConfigOptionBooleanRequest, SetSessionConfigOptionResponse, SetSessionConfigOptionSelectRequest, @@ -52,6 +56,56 @@ _CLIENT_CONNECTION_ERROR = "ClientSideConnection requires asyncio StreamWriter/StreamReader" +class _SessionUpdateTracker: + """Client proxy that tracks in-flight session updates.""" + + def __init__(self, client: Client) -> None: + self._client = client + self._session_update, self._session_update_attr, self._legacy_session_update = _resolve_handler( + client, "session_update" + ) + self._pending: dict[str, set[asyncio.Future[None]]] = {} + self._current_update: ContextVar[asyncio.Future[None] | None] = ContextVar( + "acp_current_session_update", default=None + ) + + async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None: + completed: asyncio.Future[None] = asyncio.get_running_loop().create_future() + pending = self._pending.setdefault(session_id, set()) + pending.add(completed) + token = self._current_update.set(completed) + + try: + if self._session_update is None: + raise RequestError.method_not_found(CLIENT_METHODS["session_update"]) + if self._legacy_session_update: + _warn_legacy_handler(self._client, self._session_update_attr) + notification = SessionNotification(session_id=session_id, update=update, field_meta=kwargs or None) + await self._session_update(notification) + else: + await self._session_update(session_id=session_id, update=update, **kwargs) + finally: + self._current_update.reset(token) + if not completed.done(): + completed.set_result(None) + pending.discard(completed) + if not pending: + self._pending.pop(session_id, None) + + async def wait(self, session_id: str) -> None: + # Snapshot before yielding so updates received after the response are + # not associated with this prompt. + current = self._current_update.get() + notifications = tuple( + completed for completed in self._pending.get(session_id, set()) if completed is not current + ) + if notifications: + await asyncio.gather(*(asyncio.shield(completed) for completed in notifications)) + + def __getattr__(self, name: str) -> Any: + return getattr(self._client, name) + + @final @compatible_class class ClientSideConnection: @@ -69,7 +123,9 @@ def __init__( **connection_kwargs: Any, ) -> None: client = to_client(self) if callable(to_client) else to_client - handler = build_client_router(cast(Client, client), use_unstable_protocol=use_unstable_protocol) + self._session_updates = _SessionUpdateTracker(cast(Client, client)) + handler = build_client_router(cast(Client, self._session_updates), use_unstable_protocol=use_unstable_protocol) + if isinstance(input_stream, Transport): if output_stream is not None: raise TypeError(_CLIENT_CONNECTION_ERROR) @@ -206,12 +262,18 @@ async def prompt( ], **kwargs: Any, ) -> PromptResponse: - return await request_model( - self._conn, - AGENT_METHODS["session_prompt"], - PromptRequest(prompt=prompt, session_id=session_id, field_meta=kwargs or None), - PromptResponse, - ) + try: + response = await request_model( + self._conn, + AGENT_METHODS["session_prompt"], + PromptRequest(prompt=prompt, session_id=session_id, field_meta=kwargs or None), + PromptResponse, + ) + except Exception: + await self._session_updates.wait(session_id) + raise + await self._session_updates.wait(session_id) + return response @param_model(ForkSessionRequest) async def fork_session( diff --git a/tests/test_rpc.py b/tests/test_rpc.py index 1d29465..ef4623e 100644 --- a/tests/test_rpc.py +++ b/tests/test_rpc.py @@ -34,6 +34,7 @@ ) from acp.connection import Connection from acp.core import AgentSideConnection, ClientSideConnection +from acp.exceptions import RequestError from acp.schema import ( AgentMessageChunk, AllowedOutcome, @@ -144,6 +145,154 @@ async def test_session_notifications_flow(connect, client): assert client.notifications[0].session_id == "sess" +@pytest.mark.asyncio +async def test_response_waits_for_preceding_notification(server): + notification_started = asyncio.Event() + release_notification = asyncio.Event() + + class _BlockingClient(TestClient): + async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None: + notification_started.set() + await release_notification.wait() + await super().session_update(session_id, update, **kwargs) + + client = _BlockingClient() + conn = ClientSideConnection(client, server.client_writer, server.client_reader) + request = asyncio.create_task( + conn.prompt(session_id="sess", prompt=[TextContentBlock(type="text", text="question")]) + ) + + request_message = json.loads(await server.server_reader.readline()) + notification = { + "jsonrpc": "2.0", + "method": "session/update", + "params": { + "sessionId": "sess", + "update": { + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": "answer"}, + }, + }, + } + response = {"jsonrpc": "2.0", "id": request_message["id"], "result": {"stopReason": "end_turn"}} + server.server_writer.write((json.dumps(notification) + "\n" + json.dumps(response) + "\n").encode()) + await server.server_writer.drain() + + await asyncio.wait_for(notification_started.wait(), timeout=1) + await asyncio.sleep(0) + assert not request.done() + + release_notification.set() + prompt_response = await asyncio.wait_for(request, timeout=1) + assert prompt_response.stop_reason == "end_turn" + assert len(client.notifications) == 1 + assert client.notifications[0].session_id == "sess" + await conn.close() + + +@pytest.mark.asyncio +async def test_error_response_waits_for_preceding_notification(server): + notification_started = asyncio.Event() + release_notification = asyncio.Event() + + class _BlockingClient(TestClient): + async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None: + notification_started.set() + await release_notification.wait() + + conn = ClientSideConnection(_BlockingClient(), server.client_writer, server.client_reader) + request = asyncio.create_task( + conn.prompt(session_id="sess", prompt=[TextContentBlock(type="text", text="question")]) + ) + + request_message = json.loads(await server.server_reader.readline()) + notification = { + "jsonrpc": "2.0", + "method": "session/update", + "params": { + "sessionId": "sess", + "update": { + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": "partial answer"}, + }, + }, + } + response = { + "jsonrpc": "2.0", + "id": request_message["id"], + "error": {"code": -32603, "message": "prompt failed"}, + } + server.server_writer.write((json.dumps(notification) + "\n" + json.dumps(response) + "\n").encode()) + await server.server_writer.drain() + + await asyncio.wait_for(notification_started.wait(), timeout=1) + await asyncio.sleep(0) + assert not request.done() + + release_notification.set() + with pytest.raises(RequestError, match="prompt failed"): + await asyncio.wait_for(request, timeout=1) + await conn.close() + + +@pytest.mark.asyncio +async def test_notification_can_await_nested_request(server): + notification_finished = asyncio.Event() + + class _NestedPromptClient(TestClient): + def __init__(self) -> None: + super().__init__() + self.conn: Agent | None = None + self.nested_result: PromptResponse | None = None + + def on_connect(self, conn: Agent) -> None: + self.conn = conn + + async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None: + assert self.conn is not None + self.nested_result = await self.conn.prompt( + session_id=session_id, + prompt=[TextContentBlock(type="text", text="nested question")], + ) + notification_finished.set() + + client = _NestedPromptClient() + conn = ClientSideConnection(client, server.client_writer, server.client_reader) + outer_request = asyncio.create_task( + conn.prompt(session_id="sess", prompt=[TextContentBlock(type="text", text="outer question")]) + ) + outer_message = json.loads(await server.server_reader.readline()) + + notification = { + "jsonrpc": "2.0", + "method": "session/update", + "params": { + "sessionId": "sess", + "update": { + "sessionUpdate": "agent_message_chunk", + "content": {"type": "text", "text": "answer"}, + }, + }, + } + server.server_writer.write((json.dumps(notification) + "\n").encode()) + await server.server_writer.drain() + + nested_message = json.loads(await asyncio.wait_for(server.server_reader.readline(), timeout=1)) + nested_response = {"jsonrpc": "2.0", "id": nested_message["id"], "result": {"stopReason": "end_turn"}} + server.server_writer.write((json.dumps(nested_response) + "\n").encode()) + await server.server_writer.drain() + + await asyncio.wait_for(notification_finished.wait(), timeout=1) + assert client.nested_result is not None + assert client.nested_result.stop_reason == "end_turn" + + outer_response = {"jsonrpc": "2.0", "id": outer_message["id"], "result": {"stopReason": "end_turn"}} + server.server_writer.write((json.dumps(outer_response) + "\n").encode()) + await server.server_writer.drain() + assert (await asyncio.wait_for(outer_request, timeout=1)).stop_reason == "end_turn" + await conn.close() + + @pytest.mark.asyncio async def test_on_connect_create_terminal_handle(server): class _TerminalAgent(Agent):