fix(connection): retain ordered responses across EOF

Author: hallerite

src/acp/connection.py

@@ -55,6 +55,13 @@ class StreamEvent:
     message: dict[str, Any]
 
 
+@dataclass(slots=True)
+class _RequestNotificationState:
+    start_sequence: int
+    barrier: asyncio.Future[None]
+    response_received: bool = False
+
+
 StreamObserver = Callable[[StreamEvent], Awaitable[None] | None]
 
 
@@ -81,7 +88,7 @@ def __init__(
         # 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._request_notifications: dict[int, _RequestNotificationState] = {}
         self._state = state_store or InMemoryMessageStateStore()
         self._tasks = TaskSupervisor(source="acp.Connection")
         self._tasks.add_error_handler(self._on_task_error)
@@ -126,6 +133,7 @@ async def close(self) -> None:
         await self._dispatcher.stop()
         await self._transport.close()
         await self._tasks.shutdown()
+        self._release_request_barriers()
         self._state.reject_all_outgoing(ConnectionError("Connection closed"))
 
     async def main_loop(self) -> None:
@@ -150,7 +158,11 @@ 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
+        notification_state = _RequestNotificationState(
+            start_sequence=self._notification_sequence,
+            barrier=asyncio.get_running_loop().create_future(),
+        )
+        self._request_notifications[request_id] = notification_state
         future = self._state.register_outgoing(request_id, method)
         payload = {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}
         try:
@@ -159,14 +171,18 @@ 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._request_notifications.pop(request_id, None)
             self._state.reject_outgoing(request_id, exc)
             raise
         self._notify_observers(StreamDirection.OUTGOING, payload)
         try:
+            await notification_state.barrier
             return await future
+        except asyncio.CancelledError:
+            future.cancel()
+            raise
         finally:
-            self._request_notification_starts.pop(request_id, None)
+            self._request_notifications.pop(request_id, None)
 
     async def send_notification(self, method: str, params: JsonValue | None = None) -> None:
         self._raise_if_unavailable()
@@ -204,29 +220,41 @@ async def _process_message(self, message: dict[str, Any]) -> None:
             return
         if has_id:
             request_id = message["id"]
+            notification_state = self._request_notifications.get(request_id)
+            if notification_state is None:
+                await self._handle_response(message)
+                return
             # 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
+                completion
+                for sequence, completion in self._pending_notifications.items()
+                if sequence > notification_state.start_sequence
             )
+            # Resolve the stored response before waiting. Otherwise EOF can
+            # reject a response that was already received while its preceding
+            # notification handler is still running.
+            await self._handle_response(message)
+            notification_state.response_received = True
             if preceding_notifications:
                 self._tasks.create(
-                    self._handle_response_after_notifications(message, preceding_notifications),
-                    name="acp.Connection.response",
-                    on_error=self._on_receive_error,
+                    self._release_response_after_notifications(notification_state, preceding_notifications),
+                    name="acp.Connection.response-barrier",
                 )
-            else:
-                await self._handle_response(message)
+            elif not notification_state.barrier.done():
+                notification_state.barrier.set_result(None)
 
-    async def _handle_response_after_notifications(
+    async def _release_response_after_notifications(
         self,
-        message: dict[str, Any],
+        notification_state: _RequestNotificationState,
         preceding_notifications: tuple[asyncio.Future[None], ...],
     ) -> None:
-        await asyncio.gather(*(asyncio.shield(completion) for completion in preceding_notifications))
-        await self._handle_response(message)
+        try:
+            await asyncio.gather(*(asyncio.shield(completion) for completion in preceding_notifications))
+        finally:
+            if not notification_state.barrier.done():
+                notification_state.barrier.set_result(None)
 
     def _notify_observers(self, direction: StreamDirection, message: dict[str, Any]) -> None:
         if not self._observers:
@@ -357,8 +385,16 @@ def _disconnect(self) -> None:
         if self._disconnected:
             return
         self._disconnected = True
+        self._release_request_barriers(response_received=False)
         self._state.reject_all_outgoing(ConnectionError("Connection closed"))
 
+    def _release_request_barriers(self, *, response_received: bool | None = None) -> None:
+        for state in self._request_notifications.values():
+            if response_received is not None and state.response_received is not response_received:
+                continue
+            if not state.barrier.done():
+                state.barrier.set_result(None)
+
     def _raise_if_unavailable(self) -> None:
         if self._disconnected or self._closed:
             raise ConnectionError("Connection closed")

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,
@@ -180,6 +181,79 @@ async def handler(method: str, params: Any, is_notification: bool) -> None:
     await conn.close()
 
 
+@pytest.mark.asyncio
+async def test_response_received_before_eof_waits_for_notification(server):
+    notification_started = asyncio.Event()
+    release_notification = asyncio.Event()
+
+    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()
+
+    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)
+
+    server.server_writer.close()
+    await server.server_writer.wait_closed()
+    await asyncio.sleep(0)
+    assert not request.done()
+
+    release_notification.set()
+    assert await asyncio.wait_for(request, timeout=1) == {"stopReason": "end_turn"}
+    await conn.close()
+
+
+@pytest.mark.asyncio
+async def test_error_response_waits_for_preceding_notification(server):
+    notification_started = asyncio.Event()
+    release_notification = asyncio.Event()
+
+    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()
+
+    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": "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):
     nested_result: Any = None