Skip to content

fix(connection): preserve notification response ordering #129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hallerite wants to merge 1 commit into
base: main
Choose a base branch
from
+212 −9
Open

fix(connection): preserve notification response ordering #129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 63 additions & 9 deletions src/acp/client/connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
from __future__ import annotations

import asyncio
from collections.abc import Callable
from collections.abc import Awaitable, Callable
from typing import Any, cast, final

from .._transport import Transport
from ..connection import Connection
from ..interfaces import Agent, Client
from ..meta import AGENT_METHODS
from ..meta import AGENT_METHODS, CLIENT_METHODS
from ..schema import (
AcpMcpServer,
AudioContentBlock,
Expand Down Expand Up @@ -37,6 +37,7 @@
ResourceContentBlock,
ResumeSessionRequest,
ResumeSessionResponse,
SessionNotification,
SetSessionConfigOptionBooleanRequest,
SetSessionConfigOptionResponse,
SetSessionConfigOptionSelectRequest,
Expand All @@ -52,6 +53,40 @@
_CLIENT_CONNECTION_ERROR = "ClientSideConnection requires asyncio StreamWriter/StreamReader"


class _SessionUpdateTracker:
"""Track in-flight session updates relative to each prompt."""

def __init__(self) -> None:
self._latest: dict[str, int] = {}
self._pending: dict[str, dict[int, asyncio.Future[None]]] = {}

def checkpoint(self, session_id: str) -> int:
return self._latest.get(session_id, 0)

async def handle(self, session_id: str, notification: Awaitable[Any]) -> Any:
sequence = self._latest.get(session_id, 0) + 1
self._latest[session_id] = sequence
completed: asyncio.Future[None] = asyncio.get_running_loop().create_future()
pending = self._pending.setdefault(session_id, {})
pending[sequence] = completed
try:
return await notification
finally:
if not completed.done():
completed.set_result(None)
pending.pop(sequence, None)
if not pending:
self._pending.pop(session_id, None)

async def wait(self, session_id: str, after: int) -> None:
# Snapshot before yielding so updates received after the response are
# not associated with this prompt.
pending = self._pending.get(session_id, {})
notifications = tuple(completed for sequence, completed in pending.items() if sequence > after)
if notifications:
await asyncio.gather(*(asyncio.shield(completed) for completed in notifications))


@final
@compatible_class
class ClientSideConnection:
Expand All @@ -69,7 +104,17 @@ def __init__(
**connection_kwargs: Any,
) -> None:
client = to_client(self) if callable(to_client) else to_client
handler = build_client_router(cast(Client, client), use_unstable_protocol=use_unstable_protocol)
router = build_client_router(cast(Client, client), use_unstable_protocol=use_unstable_protocol)
self._session_updates = _SessionUpdateTracker()

async def handler(method: str, params: Any, is_notification: bool) -> Any:
if is_notification and method == CLIENT_METHODS["session_update"]:
notification = SessionNotification.model_validate(params)
return await self._session_updates.handle(
notification.session_id, router(method, params, is_notification)
)
return await router(method, params, is_notification)
Comment on lines +110 to +116

@frostming frostming Aug 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change:

class _SessionUpdateTracker:
    def __init__(self, client):
        self._client = client
        self._pending = {}  # session_id -> set of unresolved futures
        # no need to track the sequence numbers, all prompt requests of the same session
        # can wait for the same set of futures.

    async def session_update(self, session_id, update):
        # tracking logic
        return await self._client.session_update(session_id, update)

    # other useful methods

    def __getattr__(self, name):
        """delegate to the internal client methods"""
        return getattr(self._client, name)

And pass the tracker to the router builder:

handler = build_client_router(self._session_updates, ...)


if isinstance(input_stream, Transport):
if output_stream is not None:
raise TypeError(_CLIENT_CONNECTION_ERROR)
Expand Down Expand Up @@ -206,12 +251,21 @@ async def prompt(
],
**kwargs: Any,
) -> PromptResponse:
return await request_model(
self._conn,
AGENT_METHODS["session_prompt"],
PromptRequest(prompt=prompt, session_id=session_id, field_meta=kwargs or None),
PromptResponse,
)
checkpoint = self._session_updates.checkpoint(session_id)
try:
response = await request_model(
self._conn,
AGENT_METHODS["session_prompt"],
PromptRequest(prompt=prompt, session_id=session_id, field_meta=kwargs or None),
PromptResponse,
)
except asyncio.CancelledError:
raise
Comment on lines +262 to +263

@frostming frostming Aug 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe a bit nitpicky, but this except branch can be deleted, because CancelledError won't be caught by the following except Exception.

except Exception:
await self._session_updates.wait(session_id, checkpoint)
raise
await self._session_updates.wait(session_id, checkpoint)
return response

@param_model(ForkSessionRequest)
async def fork_session(
Expand Down
149 changes: 149 additions & 0 deletions tests/test_rpc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -144,6 +145,154 @@ 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()

class _BlockingClient(TestClient):
async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None:
notification_started.set()
await release_notification.wait()
await super().session_update(session_id, update, **kwargs)

client = _BlockingClient()
conn = ClientSideConnection(client, server.client_writer, server.client_reader)
request = asyncio.create_task(
conn.prompt(session_id="sess", prompt=[TextContentBlock(type="text", text="question")])
)

request_message = json.loads(await server.server_reader.readline())
notification = {
"jsonrpc": "2.0",
"method": "session/update",
"params": {
"sessionId": "sess",
"update": {
"sessionUpdate": "agent_message_chunk",
"content": {"type": "text", "text": "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()
prompt_response = await asyncio.wait_for(request, timeout=1)
assert prompt_response.stop_reason == "end_turn"
assert len(client.notifications) == 1
assert client.notifications[0].session_id == "sess"
await conn.close()


@pytest.mark.asyncio
async def test_error_response_waits_for_preceding_notification(server):
notification_started = asyncio.Event()
release_notification = asyncio.Event()

class _BlockingClient(TestClient):
async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None:
notification_started.set()
await release_notification.wait()

conn = ClientSideConnection(_BlockingClient(), server.client_writer, server.client_reader)
request = asyncio.create_task(
conn.prompt(session_id="sess", prompt=[TextContentBlock(type="text", text="question")])
)

request_message = json.loads(await server.server_reader.readline())
notification = {
"jsonrpc": "2.0",
"method": "session/update",
"params": {
"sessionId": "sess",
"update": {
"sessionUpdate": "agent_message_chunk",
"content": {"type": "text", "text": "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):
notification_finished = asyncio.Event()

class _NestedPromptClient(TestClient):
def __init__(self) -> None:
super().__init__()
self.conn: Agent | None = None
self.nested_result: PromptResponse | None = None

def on_connect(self, conn: Agent) -> None:
self.conn = conn

async def session_update(self, session_id: str, update: Any, **kwargs: Any) -> None:
assert self.conn is not None
self.nested_result = await self.conn.prompt(
session_id=session_id,
prompt=[TextContentBlock(type="text", text="nested question")],
)
notification_finished.set()

client = _NestedPromptClient()
conn = ClientSideConnection(client, server.client_writer, server.client_reader)
outer_request = asyncio.create_task(
conn.prompt(session_id="sess", prompt=[TextContentBlock(type="text", text="outer question")])
)
outer_message = json.loads(await server.server_reader.readline())

notification = {
"jsonrpc": "2.0",
"method": "session/update",
"params": {
"sessionId": "sess",
"update": {
"sessionUpdate": "agent_message_chunk",
"content": {"type": "text", "text": "answer"},
},
},
}
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": {"stopReason": "end_turn"}}
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 client.nested_result is not None
assert client.nested_result.stop_reason == "end_turn"

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)).stop_reason == "end_turn"
await conn.close()


@pytest.mark.asyncio
async def test_on_connect_create_terminal_handle(server):
class _TerminalAgent(Agent):
Expand Down
Loading