fix(connection): preserve notification response ordering

Author: hallerite

src/acp/connection.py

@@ -77,6 +77,11 @@ def __init__(
     ) -> None:
         self._handler = handler
         self._next_request_id = 0
+        # Track the notification interval for each outgoing request so its
+        # response cannot overtake notifications received during that request.
+        self._notification_sequence = 0
+        self._pending_notifications: dict[int, asyncio.Future[None]] = {}
+        self._request_notification_starts: dict[int, int] = {}
         self._state = state_store or InMemoryMessageStateStore()
         self._tasks = TaskSupervisor(source="acp.Connection")
         self._tasks.add_error_handler(self._on_task_error)
@@ -145,6 +150,7 @@ async def send_request(self, method: str, params: JsonValue | None = None) -> An
         self._raise_if_unavailable()
         request_id = self._next_request_id
         self._next_request_id += 1
+        self._request_notification_starts[request_id] = self._notification_sequence
         future = self._state.register_outgoing(request_id, method)
         payload = {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}
         try:
@@ -153,10 +159,14 @@ async def send_request(self, method: str, params: JsonValue | None = None) -> An
             # A synchronous send failure (e.g. HTTP POST rejected before any
             # JSON-RPC response exists) must reject the correlated future so the
             # caller gets a real, attributable error.
+            self._request_notification_starts.pop(request_id, None)
             self._state.reject_outgoing(request_id, exc)
             raise
         self._notify_observers(StreamDirection.OUTGOING, payload)
-        return await future
+        try:
+            return await future
+        finally:
+            self._request_notification_starts.pop(request_id, None)
 
     async def send_notification(self, method: str, params: JsonValue | None = None) -> None:
         self._raise_if_unavailable()
@@ -185,10 +195,38 @@ async def _process_message(self, message: dict[str, Any]) -> None:
             await self._queue.publish(RpcTask(RpcTaskKind.REQUEST, message))
             return
         if method is not None and not has_id:
-            await self._queue.publish(RpcTask(RpcTaskKind.NOTIFICATION, message))
+            self._notification_sequence += 1
+            sequence = self._notification_sequence
+            completion = asyncio.get_running_loop().create_future()
+            self._pending_notifications[sequence] = completion
+            completion.add_done_callback(lambda _: self._pending_notifications.pop(sequence, None))
+            await self._queue.publish(RpcTask(RpcTaskKind.NOTIFICATION, message, completion))
             return
         if has_id:
-            await self._handle_response(message)
+            request_id = message["id"]
+            # Excluding notifications received before this request began keeps
+            # notification handlers free to make nested requests without those
+            # responses waiting on the handler that issued them.
+            start_sequence = self._request_notification_starts.get(request_id, self._notification_sequence)
+            preceding_notifications = tuple(
+                completion for sequence, completion in self._pending_notifications.items() if sequence > start_sequence
+            )
+            if preceding_notifications:
+                self._tasks.create(
+                    self._handle_response_after_notifications(message, preceding_notifications),
+                    name="acp.Connection.response",
+                    on_error=self._on_receive_error,
+                )
+            else:
+                await self._handle_response(message)
+
+    async def _handle_response_after_notifications(
+        self,
+        message: dict[str, Any],
+        preceding_notifications: tuple[asyncio.Future[None], ...],
+    ) -> None:
+        await asyncio.gather(*(asyncio.shield(completion) for completion in preceding_notifications))
+        await self._handle_response(message)
 
     def _notify_observers(self, direction: StreamDirection, message: dict[str, Any]) -> None:
         if not self._observers:

src/acp/task/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any
@@ -16,6 +17,7 @@ class RpcTaskKind(Enum):
 class RpcTask:
     kind: RpcTaskKind
     message: dict[str, Any]
+    completion: asyncio.Future[None] | None = None
 
 
 from .dispatcher import (  # noqa: E402

src/acp/task/dispatcher.py

@@ -5,7 +5,7 @@
 from contextlib import suppress
 from typing import Any, Protocol
 
-from . import RpcTaskKind
+from . import RpcTask, RpcTaskKind
 from .queue import MessageQueue
 from .state import MessageStateStore
 from .supervisor import TaskSupervisor
@@ -60,7 +60,7 @@ async def _run(self) -> None:
                     if task.kind is RpcTaskKind.REQUEST:
                         await self._dispatch_request(task.message)
                     else:
-                        await self._dispatch_notification(task.message)
+                        await self._dispatch_notification(task)
                 finally:
                     self._queue.task_done()
         except asyncio.CancelledError:
@@ -87,8 +87,12 @@ async def runner() -> None:
 
         self._supervisor.create(runner(), name="acp.Dispatcher.request")
 
-    async def _dispatch_notification(self, message: dict[str, Any]) -> None:
+    async def _dispatch_notification(self, task: RpcTask) -> None:
         async def runner() -> None:
-            await self._notification_runner(message)
+            try:
+                await self._notification_runner(task.message)
+            finally:
+                if task.completion is not None and not task.completion.done():
+                    task.completion.set_result(None)
 
         self._supervisor.create(runner(), name="acp.Dispatcher.notification")

tests/test_rpc.py

@@ -144,6 +144,83 @@ 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()
+    notifications: list[Any] = []
+
+    async def handler(method: str, params: Any, is_notification: bool) -> None:
+        assert method == "session/update"
+        assert is_notification
+        notification_started.set()
+        await release_notification.wait()
+        notifications.append(params)
+
+    conn = Connection(handler, server.client_writer, server.client_reader)
+    request = asyncio.create_task(conn.send_request("session/prompt", {"sessionId": "sess"}))
+
+    request_message = json.loads(await server.server_reader.readline())
+    notification = {
+        "jsonrpc": "2.0",
+        "method": "session/update",
+        "params": {"sessionId": "sess", "update": "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()
+    assert await asyncio.wait_for(request, timeout=1) == {"stopReason": "end_turn"}
+    assert notifications == [notification["params"]]
+    await conn.close()
+
+
+@pytest.mark.asyncio
+async def test_notification_can_await_nested_request(server):
+    nested_result: Any = None
+    notification_finished = asyncio.Event()
+    conn: Connection | None = None
+
+    async def handler(method: str, params: Any, is_notification: bool) -> None:
+        nonlocal nested_result
+        assert conn is not None
+        assert method == "session/update"
+        assert is_notification
+        nested_result = await conn.send_request("nested/request", params)
+        notification_finished.set()
+
+    conn = Connection(handler, server.client_writer, server.client_reader)
+    outer_request = asyncio.create_task(conn.send_request("session/prompt", {"sessionId": "sess"}))
+    outer_message = json.loads(await server.server_reader.readline())
+
+    notification = {
+        "jsonrpc": "2.0",
+        "method": "session/update",
+        "params": {"sessionId": "sess"},
+    }
+    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": {"ok": True}}
+    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 nested_result == {"ok": True}
+
+    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) == {"stopReason": "end_turn"}
+    await conn.close()
+
+
 @pytest.mark.asyncio
 async def test_on_connect_create_terminal_handle(server):
     class _TerminalAgent(Agent):