fix(client): preserve session update ordering (#129)

Author: hallerite

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(

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):