fix(connection): preserve notification response ordering

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")

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

@@ -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):