diff --git a/src/acp/connection.py b/src/acp/connection.py index cfd7b5c..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] @@ -77,6 +84,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_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) @@ -121,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: @@ -145,6 +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 + 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: @@ -153,10 +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_notifications.pop(request_id, None) self._state.reject_outgoing(request_id, exc) raise self._notify_observers(StreamDirection.OUTGOING, payload) - return await future + try: + await notification_state.barrier + return await future + except asyncio.CancelledError: + future.cancel() + raise + finally: + self._request_notifications.pop(request_id, None) async def send_notification(self, method: str, params: JsonValue | None = None) -> None: self._raise_if_unavailable() @@ -185,10 +211,50 @@ 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: + 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. + preceding_notifications = tuple( + 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._release_response_after_notifications(notification_state, preceding_notifications), + name="acp.Connection.response-barrier", + ) + elif not notification_state.barrier.done(): + notification_state.barrier.set_result(None) + + async def _release_response_after_notifications( + self, + notification_state: _RequestNotificationState, + preceding_notifications: tuple[asyncio.Future[None], ...], + ) -> None: + 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: @@ -319,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/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..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, @@ -144,6 +145,121 @@ 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_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 + 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):