From 921dc9aeb0311daa205e5f1bc4e1f2000e0b1ebf Mon Sep 17 00:00:00 2001 From: hallerite Date: Thu, 6 Aug 2026 01:03:59 +0200 Subject: [PATCH 1/2] fix(connection): preserve notification response ordering --- src/acp/connection.py | 44 ++++++++++++++++++++-- src/acp/task/__init__.py | 2 + src/acp/task/dispatcher.py | 12 ++++-- tests/test_rpc.py | 77 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 7 deletions(-) diff --git a/src/acp/connection.py b/src/acp/connection.py index cfd7b5c..97a9696 100644 --- a/src/acp/connection.py +++ b/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: diff --git a/src/acp/task/__init__.py b/src/acp/task/__init__.py index 2896fbf..2dd2951 100644 --- a/src/acp/task/__init__.py +++ b/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 diff --git a/src/acp/task/dispatcher.py b/src/acp/task/dispatcher.py index e8c5e76..38ae0f6 100644 --- a/src/acp/task/dispatcher.py +++ b/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") diff --git a/tests/test_rpc.py b/tests/test_rpc.py index 5e8a917..fe0b96f 100644 --- a/tests/test_rpc.py +++ b/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): From a966b23568aa95f7ebf582cf3c9daaa9d9ea65f3 Mon Sep 17 00:00:00 2001 From: hallerite Date: Thu, 6 Aug 2026 12:56:13 +0200 Subject: [PATCH 2/2] fix(connection): retain ordered responses across EOF --- src/acp/connection.py | 66 +++++++++++++++++++++++++++++++++---------- tests/test_rpc.py | 39 +++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 15 deletions(-) diff --git a/src/acp/connection.py b/src/acp/connection.py index 97a9696..3000eba 100644 --- a/src/acp/connection.py +++ b/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") diff --git a/tests/test_rpc.py b/tests/test_rpc.py index fe0b96f..e4925b0 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, @@ -180,6 +181,44 @@ async def handler(method: str, params: Any, is_notification: bool) -> None: 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