fix(client): preserve session update ordering

Author: hallerite

src/acp/client/connection.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from typing import Any, cast, final
 
 from .._transport import Transport
 from ..connection import Connection
 from ..interfaces import Agent, Client
-from ..meta import AGENT_METHODS
+from ..meta import AGENT_METHODS, CLIENT_METHODS
 from ..schema import (
     AcpMcpServer,
     AudioContentBlock,
@@ -37,6 +37,7 @@
     ResourceContentBlock,
     ResumeSessionRequest,
     ResumeSessionResponse,
+    SessionNotification,
     SetSessionConfigOptionBooleanRequest,
     SetSessionConfigOptionResponse,
     SetSessionConfigOptionSelectRequest,
@@ -52,6 +53,40 @@
 _CLIENT_CONNECTION_ERROR = "ClientSideConnection requires asyncio StreamWriter/StreamReader"
 
 
+class _SessionUpdateTracker:
+    """Track in-flight session updates relative to each prompt."""
+
+    def __init__(self) -> None:
+        self._latest: dict[str, int] = {}
+        self._pending: dict[str, dict[int, asyncio.Future[None]]] = {}
+
+    def checkpoint(self, session_id: str) -> int:
+        return self._latest.get(session_id, 0)
+
+    async def handle(self, session_id: str, notification: Awaitable[Any]) -> Any:
+        sequence = self._latest.get(session_id, 0) + 1
+        self._latest[session_id] = sequence
+        completed: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+        pending = self._pending.setdefault(session_id, {})
+        pending[sequence] = completed
+        try:
+            return await notification
+        finally:
+            if not completed.done():
+                completed.set_result(None)
+            pending.pop(sequence, None)
+            if not pending:
+                self._pending.pop(session_id, None)
+
+    async def wait(self, session_id: str, after: int) -> None:
+        # Snapshot before yielding so updates received after the response are
+        # not associated with this prompt.
+        pending = self._pending.get(session_id, {})
+        notifications = tuple(completed for sequence, completed in pending.items() if sequence > after)
+        if notifications:
+            await asyncio.gather(*(asyncio.shield(completed) for completed in notifications))
+
+
 @final
 @compatible_class
 class ClientSideConnection:
@@ -69,7 +104,17 @@ 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)
+        router = build_client_router(cast(Client, client), use_unstable_protocol=use_unstable_protocol)
+        self._session_updates = _SessionUpdateTracker()
+
+        async def handler(method: str, params: Any, is_notification: bool) -> Any:
+            if is_notification and method == CLIENT_METHODS["session_update"]:
+                notification = SessionNotification.model_validate(params)
+                return await self._session_updates.handle(
+                    notification.session_id, router(method, params, is_notification)
+                )
+            return await router(method, params, is_notification)
+
         if isinstance(input_stream, Transport):
             if output_stream is not None:
                 raise TypeError(_CLIENT_CONNECTION_ERROR)
@@ -206,12 +251,21 @@ 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,
-        )
+        checkpoint = self._session_updates.checkpoint(session_id)
+        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 asyncio.CancelledError:
+            raise
+        except Exception:
+            await self._session_updates.wait(session_id, checkpoint)
+            raise
+        await self._session_updates.wait(session_id, checkpoint)
+        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):