diff --git a/src/acp/_jsonrpc.py b/src/acp/_jsonrpc.py index 078c5621..e9532dd2 100644 --- a/src/acp/_jsonrpc.py +++ b/src/acp/_jsonrpc.py @@ -5,13 +5,20 @@ from .exceptions import RequestError +JsonRpcMessage = dict[str, Any] | list[Any] -def decode_message(payload: str | bytes) -> dict[str, Any] | None: - """Decode a message, returning None for unsupported non-empty batches.""" + +def decode_message(payload: str | bytes) -> JsonRpcMessage | None: + """Decode a JSON-RPC message, returning a non-empty batch as a list.""" try: message = json.loads(payload) except (json.JSONDecodeError, UnicodeDecodeError) as exc: raise RequestError.parse_error() from exc + if isinstance(message, list): + # An empty batch is not a batch: the spec answers it with a single -32600 object. + if not message: + raise RequestError.invalid_request() + return message return require_message_object(message) diff --git a/src/acp/_sse.py b/src/acp/_sse.py index 2f67e3e2..e4cfd110 100644 --- a/src/acp/_sse.py +++ b/src/acp/_sse.py @@ -9,9 +9,9 @@ import codecs import json -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING -from ._jsonrpc import decode_message +from ._jsonrpc import JsonRpcMessage, decode_message from .exceptions import RequestError if TYPE_CHECKING: @@ -24,7 +24,7 @@ ] -def serialize_sse_event(message: dict[str, Any]) -> bytes: +def serialize_sse_event(message: JsonRpcMessage) -> bytes: """Serialize a JSON-RPC message as an SSE ``data:`` event. The payload is JSON-encoded on a single line and terminated by a blank line, @@ -41,8 +41,8 @@ def serialize_sse_keepalive() -> bytes: async def _decode_event( data_lines: list[str], on_error: Callable[[RequestError], Awaitable[None]] | None -) -> dict[str, Any] | None: - """Decode buffered ``data:`` lines into a JSON object, or None to skip.""" +) -> JsonRpcMessage | None: + """Decode buffered ``data:`` lines into a JSON object or batch, or None to skip.""" payload = "\n".join(data_lines) if not payload: return None @@ -65,11 +65,12 @@ def _append_field(line: str, data_lines: list[str]) -> None: async def parse_sse_stream( chunks: AsyncIterator[bytes], *, on_error: Callable[[RequestError], Awaitable[None]] | None = None -) -> AsyncIterator[dict[str, Any]]: +) -> AsyncIterator[JsonRpcMessage]: """Parse an SSE byte stream, yielding decoded JSON-RPC ``data:`` payloads. Comments (lines starting with ``:``) and non-``data`` fields are ignored. - Invalid JSON and non-object JSON payloads are skipped. + Invalid JSON and other non-message JSON payloads are skipped; a JSON array + is delivered as a batch. Multi-line ``data:`` fields are concatenated with newlines per the spec. A blank line dispatches the buffered event. """ diff --git a/src/acp/_transport.py b/src/acp/_transport.py index 2e2d113a..96741491 100644 --- a/src/acp/_transport.py +++ b/src/acp/_transport.py @@ -11,7 +11,7 @@ import contextlib from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable -from ._jsonrpc import decode_message, error_response +from ._jsonrpc import JsonRpcMessage, decode_message, error_response from .exceptions import RequestError if TYPE_CHECKING: @@ -36,9 +36,9 @@ class Transport(Protocol): response exists); callers correlate such failures with pending requests. """ - async def send(self, message: dict[str, Any]) -> None: ... + async def send(self, message: dict[str, Any] | list[Any]) -> None: ... - async def receive(self) -> dict[str, Any] | None: ... + async def receive(self) -> JsonRpcMessage | None: ... async def close(self) -> None: ... @@ -58,10 +58,10 @@ def __init__( self._receive_timeout = receive_timeout self._error_send_failed = False - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: await self._sender.send(message) - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> JsonRpcMessage | None: while True: line = await self._read_line() if not line: @@ -109,19 +109,19 @@ class _MemoryTransport: def __init__( self, - outbox: asyncio.Queue[dict[str, Any] | None], - inbox: asyncio.Queue[dict[str, Any] | None], + outbox: asyncio.Queue[JsonRpcMessage | None], + inbox: asyncio.Queue[JsonRpcMessage | None], ) -> None: self._outbox = outbox self._inbox = inbox self._closed = False - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: if self._closed: raise ConnectionError("Transport closed") - await self._outbox.put(dict(message)) + await self._outbox.put(dict(message) if isinstance(message, dict) else list(message)) - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> JsonRpcMessage | None: return await self._inbox.get() async def close(self) -> None: @@ -138,8 +138,8 @@ def memory_transport_pair() -> tuple[Transport, Transport]: A message ``send`` on one end becomes available via ``receive`` on the other. Closing an end enqueues an EOF (``None``) for its peer. """ - a_to_b: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() - b_to_a: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() + a_to_b: asyncio.Queue[JsonRpcMessage | None] = asyncio.Queue() + b_to_a: asyncio.Queue[JsonRpcMessage | None] = asyncio.Queue() left = _MemoryTransport(outbox=a_to_b, inbox=b_to_a) right = _MemoryTransport(outbox=b_to_a, inbox=a_to_b) return left, right diff --git a/src/acp/connection.py b/src/acp/connection.py index 84d009cf..e5675331 100644 --- a/src/acp/connection.py +++ b/src/acp/connection.py @@ -12,6 +12,7 @@ from pydantic import BaseModel, ValidationError +from ._jsonrpc import error_response as _error_response from ._transport import NdjsonTransport, Transport from .exceptions import RequestError from .meta import PROTOCOL_METHODS @@ -33,7 +34,7 @@ class StreamDirection(str, Enum): @dataclass(frozen=True, slots=True) class StreamEvent: direction: StreamDirection - message: dict[str, Any] + message: dict[str, Any] | list[Any] StreamObserver = Callable[[StreamEvent], Awaitable[None] | None] @@ -41,6 +42,11 @@ class StreamEvent: _CANCEL_REQUEST_METHOD = PROTOCOL_METHODS["cancel_request"] +async def _answer_invalid_entry(payload: dict[str, Any]) -> dict[str, Any]: + """Answer a batch entry that is not a valid request without dispatching it.""" + return payload + + def _is_request_id(value: Any) -> bool: """Whether ``value`` is a JSON-RPC ``RequestId``: ``null``, an integer or a string.""" return value is None or isinstance(value, str) or (isinstance(value, int) and not isinstance(value, bool)) @@ -148,6 +154,9 @@ async def _receive_loop(self) -> None: message = await self._transport.receive() if message is None: break + if isinstance(message, list): + self._process_batch(message) + continue self._notify_observers(StreamDirection.INCOMING, message) self._process_message(message) except asyncio.CancelledError: @@ -177,7 +186,56 @@ def _process_message(self, message: dict[str, Any]) -> None: if has_id: # this is a response, {"id", "result" | "error"} self._handle_response(message) - def _notify_observers(self, direction: StreamDirection, message: dict[str, Any]) -> None: + def _process_batch(self, batch: list[Any]) -> None: + # JSON-RPC 2.0 batches: entries are dispatched as if they had arrived one per frame, but + # every reply is collected and sent as a single array in the order the entries appeared. + # A batch with no requests (only notifications, or entries that answer nothing such as + # responses) produces no output at all. An empty array never reaches here: ``decode_message`` + # answers it with a single ``-32600`` object on every transport that decodes frames. + replies: list[asyncio.Task[Any]] = [] + for item in batch: + entry = _error_response(RequestError.invalid_request()) + if isinstance(item, dict): + self._notify_observers(StreamDirection.INCOMING, item) + method = item.get("method") + if method is None: + if "id" in item: + # A response to one of our own requests never joins the reply array. + self._handle_response(item) + continue + elif "id" not in item: + if method == _CANCEL_REQUEST_METHOD: + self._cancel_incoming(item.get("params")) + else: + self._tasks.create(self._run_notification(item), name="acp.Connection.notification") + continue + else: + entry = item + # Non-object entries keep the null-id error response built above, as the spec requires. + if "method" in entry: + handler = self._tasks.create(self._execute_request(entry), name="acp.Connection.request") + self._track_incoming(entry["id"], handler) + else: + handler = self._tasks.create(_answer_invalid_entry(entry), name="acp.Connection.batch_entry") + replies.append(self._tasks.create(self._batch_reply(handler, entry), name="acp.Connection.batch_reply")) + if replies: + self._tasks.create(self._send_batch(replies), name="acp.Connection.batch_response") + + async def _batch_reply(self, handler: asyncio.Task[Any], entry: dict[str, Any]) -> dict[str, Any]: + # Mirrors ``_run_request``: ``asyncio.wait`` keeps a handler cancelled by + # ``$/cancel_request`` from aborting the reply for the whole batch. + await asyncio.wait({handler}) + if not handler.cancelled(): + return handler.result() + return {"jsonrpc": "2.0", "id": entry.get("id"), "error": RequestError.request_cancelled().to_error_obj()} + + async def _send_batch(self, replies: list[asyncio.Task[Any]]) -> None: + # ``gather`` preserves the order of the entries, so the array mirrors the request batch. + payload = list(await asyncio.gather(*replies)) + await self._transport.send(payload) + self._notify_observers(StreamDirection.OUTGOING, payload) + + def _notify_observers(self, direction: StreamDirection, message: dict[str, Any] | list[Any]) -> None: if not self._observers: return snapshot = copy.deepcopy(message) diff --git a/src/acp/experimental/v2/_initialization.py b/src/acp/experimental/v2/_initialization.py index 5ae6efb7..e8b6010c 100644 --- a/src/acp/experimental/v2/_initialization.py +++ b/src/acp/experimental/v2/_initialization.py @@ -24,11 +24,11 @@ def phase(self) -> InitializationPhase: def begin(self, request: schema.InitializeRequest) -> None: if self._phase != "uninitialized": raise RequestError.invalid_request({"details": "ACP v2 connections may only be initialized once"}) - if request.protocol_version != PROTOCOL_VERSION: - raise RequestError.invalid_params({ - "expectedProtocolVersion": PROTOCOL_VERSION, - "receivedProtocolVersion": request.protocol_version, - }) + # A request for any version other than PROTOCOL_VERSION is answered by the response, not + # rejected: initialization requires an unsupported request to succeed with the newest + # version this runtime supports (TCK ACP-INIT-003/201/202). The requested value is + # therefore validated only by the generated request model; `complete` still enforces that + # the response advertises PROTOCOL_VERSION. self._phase = "initializing" def complete(self, response: schema.InitializeResponse) -> None: diff --git a/src/acp/experimental/v2/_router.py b/src/acp/experimental/v2/_router.py index 4afb0e7e..988b9548 100644 --- a/src/acp/experimental/v2/_router.py +++ b/src/acp/experimental/v2/_router.py @@ -37,7 +37,11 @@ async def handle_request(self, spec: RequestSpec, params: Any) -> Any: handler = self._handler(spec.handler) if handler is None: raise RequestError.method_not_found(spec.method) - request = spec.request.validate_python(params) + # ``params`` is optional on every ACP request envelope, so an omitted one means "no + # arguments" and validates like ``{}``. A request whose fields are required still fails + # validation, so this only rescues methods such as ``session/list`` whose fields are all + # optional (TCK ACP-BATCH-204/205/JSONRPC-001 send ``session/list`` with no params). + request = spec.request.validate_python({} if params is None else params) response = await handler(**model_to_kwargs(request, type(request))) if response is None and spec.empty_response: response = {} diff --git a/src/acp/http/client.py b/src/acp/http/client.py index 5087bcf9..4d026803 100644 --- a/src/acp/http/client.py +++ b/src/acp/http/client.py @@ -23,7 +23,7 @@ import contextlib from typing import TYPE_CHECKING, Any -from .._jsonrpc import error_response +from .._jsonrpc import JsonRpcMessage, error_response from .._sse import parse_sse_stream from ..exceptions import RequestError from .protocol import ( @@ -88,9 +88,13 @@ def __init__( # -- Transport protocol ------------------------------------------------- - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: if self._closed: raise ConnectionError("Transport closed") + if isinstance(message, list): + # A Streamable HTTP POST carries a single JSON-RPC message and the server answers a + # batch root with 501, so a batch can never be delivered over this transport. + raise ConnectionError("Batch requests are not supported over Streamable HTTP") if is_initialize_request(message): await self._send_initialize(message) return @@ -105,7 +109,7 @@ async def send(self, message: dict[str, Any]) -> None: self._pending_loads.pop(key, None) raise - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> JsonRpcMessage | None: item = await self._inbox.get() if item is _EOF: return None @@ -221,7 +225,12 @@ async def _send_protocol_error(self, error: RequestError) -> None: with contextlib.suppress(ConnectionError, httpx2.HTTPError, AcpHttpStatusError): await self.send(error_response(error)) - def _handle_incoming(self, message: dict[str, Any]) -> None: + def _handle_incoming(self, message: dict[str, Any] | list[Any]) -> None: + if isinstance(message, list): + for entry in message: + if isinstance(entry, dict): + self._handle_incoming(entry) + return # Load and resume responses may omit sessionId; the ID is in the request. if is_response_message(message): key = message_id_key(message.get("id")) diff --git a/src/acp/task/sender.py b/src/acp/task/sender.py index 613535bf..843af775 100644 --- a/src/acp/task/sender.py +++ b/src/acp/task/sender.py @@ -25,7 +25,7 @@ def __init__(self, writer: asyncio.StreamWriter, supervisor: TaskSupervisor) -> self._closed = False self._task = supervisor.create(self._loop(), name="acp.Sender.loop", on_error=self._on_error) - async def send(self, payload: dict[str, Any]) -> None: + async def send(self, payload: dict[str, Any] | list[Any]) -> None: data = (json.dumps(payload, separators=(",", ":")) + "\n").encode("utf-8") future: asyncio.Future[None] = asyncio.get_running_loop().create_future() await self._queue.put(_PendingSend(data, future)) diff --git a/src/acp/ws/client.py b/src/acp/ws/client.py index 2ac2da7a..51ad2699 100644 --- a/src/acp/ws/client.py +++ b/src/acp/ws/client.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any from .._cookies import MemoryAcpCookieStore -from .._jsonrpc import decode_message, error_response +from .._jsonrpc import JsonRpcMessage, decode_message, error_response from ..exceptions import RequestError try: @@ -40,12 +40,12 @@ def __init__(self, connection: Any) -> None: self._ws = connection self._closed = False - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: if self._closed: raise ConnectionError("Transport closed") await self._ws.send(json.dumps(message, separators=(",", ":"))) - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> JsonRpcMessage | None: while True: try: frame = await self._ws.recv() diff --git a/src/acp/ws/server.py b/src/acp/ws/server.py index 0ae5dcb2..9eb7acbe 100644 --- a/src/acp/ws/server.py +++ b/src/acp/ws/server.py @@ -9,7 +9,7 @@ from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState -from .._jsonrpc import decode_message, error_response +from .._jsonrpc import JsonRpcMessage, decode_message, error_response from ..agent.connection import AgentSideConnection from ..exceptions import RequestError from ..http.protocol import CONNECTION_ID_HEADER @@ -28,12 +28,12 @@ def __init__(self, websocket: WebSocket) -> None: self._ws = websocket self._warned_binary = False - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: if self._ws.client_state == WebSocketState.DISCONNECTED: raise ConnectionError("Transport closed") await self._ws.send_json(message) - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> JsonRpcMessage | None: while self._ws.client_state != WebSocketState.DISCONNECTED: event = await self._ws.receive() if event["type"] == "websocket.disconnect": diff --git a/tests/http/test_http_client.py b/tests/http/test_http_client.py index 38ddb78b..a0ad63c7 100644 --- a/tests/http/test_http_client.py +++ b/tests/http/test_http_client.py @@ -116,6 +116,21 @@ def handler(request: httpx2.Request) -> httpx2.Response: await client.aclose() +@pytest.mark.asyncio +async def test_batch_send_is_rejected_before_reaching_the_wire() -> None: + # The ACP HTTP server answers a batch POST root with 501, so the client transport refuses one + # front instead of opening a connection and failing per entry. + server = FakeServer() + transport, client = _make_transport(server) + try: + with pytest.raises(ConnectionError, match="not supported over Streamable HTTP"): + await transport.send([{"jsonrpc": "2.0", "id": 1, "method": "session/list"}]) + assert server.posts == [] + finally: + await transport.close() + await client.aclose() + + @pytest.mark.asyncio async def test_connection_scoped_sse_delivers_new_session_result() -> None: server = FakeServer() diff --git a/tests/http/test_sse.py b/tests/http/test_sse.py index cc1e1704..6449bccd 100644 --- a/tests/http/test_sse.py +++ b/tests/http/test_sse.py @@ -75,8 +75,15 @@ async def test_parse_skips_invalid_json() -> None: @pytest.mark.asyncio -@pytest.mark.parametrize("payload", [b"[]", b'[ {"id":2} ]', b'"text"', b"42", b"true", b"false", b"null"]) +@pytest.mark.parametrize("payload", [b"[]", b'"text"', b"42", b"true", b"false", b"null"]) async def test_parse_skips_non_object_json(payload: bytes) -> None: stream = _aiter([b"data: " + payload + b'\n\ndata: {"id":1}\n\n']) events = [event async for event in parse_sse_stream(stream)] assert events == [{"id": 1}] + + +@pytest.mark.asyncio +async def test_parse_yields_a_batch_array() -> None: + stream = _aiter([b'data: [ {"id":2} ]\n\ndata: {"id":1}\n\n']) + events = [event async for event in parse_sse_stream(stream)] + assert events == [[{"id": 2}], {"id": 1}] diff --git a/tests/http/test_transport_errors.py b/tests/http/test_transport_errors.py index da2d47ed..31461bc3 100644 --- a/tests/http/test_transport_errors.py +++ b/tests/http/test_transport_errors.py @@ -131,3 +131,26 @@ async def test_ignored_payload_preserves_requests_without_reply(peer: _Peer, pay await _complete(peer, pending, request) next_pending, next_request = await _request(peer) await _complete(peer, next_pending, next_request) + + +@pytest.mark.asyncio +async def test_batch_root_is_answered_as_one_array(peer: _Peer) -> None: + # A JSON-RPC batch arrives as a single frame; both replies come back in one array, and the + # response entry inside the batch resolves the request we sent instead of being echoed. + pending, request = await _request(peer) + await peer.send( + json.dumps([ + {"jsonrpc": "2.0", "id": "one", "method": "probe", "params": {"n": 1}}, + {"jsonrpc": "2.0", "id": request["id"], "result": "answered"}, + {"jsonrpc": "2.0", "id": "two", "method": "probe", "params": {"n": 2}}, + ]) + ) + + assert await asyncio.wait_for(pending, timeout=1) == "answered" + # SSE carries one message per frame, so the HTTP transport fans the array out again while + # keeping each response on its own route; stdio and WebSocket keep the single array frame. + replies: list[dict[str, Any]] = [] + while len(replies) < 2: + received = await asyncio.wait_for(peer.receive(), timeout=1) + replies.extend(received if isinstance(received, list) else [received]) + assert [reply["id"] for reply in replies] == ["one", "two"] diff --git a/tests/test_jsonrpc_batch.py b/tests/test_jsonrpc_batch.py new file mode 100644 index 00000000..df6de184 --- /dev/null +++ b/tests/test_jsonrpc_batch.py @@ -0,0 +1,212 @@ +"""Wire tests for JSON-RPC 2.0 batches over the newline-delimited transport. + +`Connection` exchanges decoded messages with its transport, so batch handling lives in +`Connection._process_batch`: every request entry is dispatched like a single request, and the +replies are collected into one ordered array. A batch that carries no request stays silent. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +from collections.abc import AsyncIterator +from typing import Any + +import pytest + +from acp._jsonrpc import decode_message +from acp._transport import memory_transport_pair +from acp.connection import Connection +from acp.exceptions import RequestError +from acp.meta import PROTOCOL_METHODS + + +def _conn(server, handler) -> Connection: + return Connection(handler, server.server_writer, server.server_reader) + + +async def _write(server, payload: Any) -> None: + server.client_writer.write((json.dumps(payload) + "\n").encode()) + await server.client_writer.drain() + + +async def _read(server) -> Any: + line = await asyncio.wait_for(server.client_reader.readline(), timeout=1) + return json.loads(line) + + +@contextlib.asynccontextmanager +async def _memory(handler) -> AsyncIterator[tuple[Connection, Any]]: + peer, transport = memory_transport_pair() + conn = Connection(handler, transport) + try: + yield conn, peer + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_batch_requests_return_one_ordered_array(server): + async def handler(method: str, params: Any, is_notification: bool) -> Any: + if params == {"slow": True}: + await asyncio.sleep(0.05) + return {"echo": params} + + conn = _conn(server, handler) + try: + await _write( + server, + [ + {"jsonrpc": "2.0", "id": 1, "method": "probe", "params": {"slow": True}}, + {"jsonrpc": "2.0", "id": "two", "method": "probe", "params": {"n": 2}}, + ], + ) + + replies = await _read(server) + + assert isinstance(replies, list), "batch replies must be a single array frame" + assert [reply["id"] for reply in replies] == [1, "two"] + assert [reply["result"] for reply in replies] == [{"echo": {"slow": True}}, {"echo": {"n": 2}}] + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_notification_only_batch_is_silent(server): + seen: list[Any] = [] + + async def handler(method: str, params: Any, is_notification: bool) -> None: + if is_notification: + seen.append(params) + + conn = _conn(server, handler) + try: + await _write( + server, + [ + {"jsonrpc": "2.0", "method": "note", "params": {"n": 1}}, + {"jsonrpc": "2.0", "method": "note", "params": {"n": 2}}, + ], + ) + # A follow-up request proves the batch produced no frame at all, rather than a slow one. + await _write(server, {"jsonrpc": "2.0", "id": 7, "method": "after"}) + + assert await _read(server) == {"jsonrpc": "2.0", "id": 7, "result": None} + assert seen == [{"n": 1}, {"n": 2}] + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_mixed_batch_answers_invalid_sibling_in_place(server): + async def handler(method: str, params: Any, is_notification: bool) -> Any: + return {"ok": True} + + conn = _conn(server, handler) + try: + await _write(server, [1, {"jsonrpc": "2.0", "id": 2, "method": "probe"}]) + + replies = await _read(server) + + assert [reply.get("id") for reply in replies] == [None, 2] + assert replies[0]["error"]["code"] == -32600 + assert replies[1]["result"] == {"ok": True} + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_empty_batch_is_a_single_invalid_request(server): + conn = _conn(server, lambda method, params, is_notification: None) + try: + await _write(server, []) + + reply = await _read(server) + + assert isinstance(reply, dict) + assert (reply["id"], reply["error"]["code"]) == (None, -32600) + with pytest.raises(RequestError) as exc_info: + decode_message("[]") + assert isinstance(exc_info.value, RequestError) + assert exc_info.value.code == -32600 + finally: + await conn.close() + + +@pytest.mark.asyncio +async def test_inbound_response_array_resolves_and_is_not_echoed(): + async def handler(method: str, params: Any, is_notification: bool) -> Any: + return None + + async with _memory(handler) as (conn, peer): + pending = asyncio.create_task(conn.send_request("probe")) + request_id = (await asyncio.wait_for(peer.receive(), timeout=1))["id"] + + await peer.send([ + {"jsonrpc": "2.0", "id": request_id, "result": {"answered": True}}, + {"jsonrpc": "2.0", "id": 9, "method": "probe"}, + ]) + + assert await asyncio.wait_for(pending, timeout=1) == {"answered": True} + assert await asyncio.wait_for(peer.receive(), timeout=1) == [{"jsonrpc": "2.0", "id": 9, "result": None}] + + +@pytest.mark.asyncio +async def test_batch_handler_can_reenter_the_connection(): + async def handler(method: str, params: Any, is_notification: bool) -> Any: + # Runs inside ``_execute_request``'s own task, so the receive loop stays free to route the + # peer's answer back into this coroutine while the batch itself is still unanswered. + answer = await conn.send_request("peer/question", {"n": params["n"]}) + return {"n": params["n"], "peer": answer} + + async with _memory(handler) as (conn, peer): + await peer.send([ + {"jsonrpc": "2.0", "id": 1, "method": "probe", "params": {"n": 1}}, + {"jsonrpc": "2.0", "id": 2, "method": "probe", "params": {"n": 2}}, + ]) + + for _ in range(2): + nested = await asyncio.wait_for(peer.receive(), timeout=1) + assert nested["method"] == "peer/question" + await peer.send({"jsonrpc": "2.0", "id": nested["id"], "result": {"answered": nested["params"]["n"]}}) + + replies = await asyncio.wait_for(peer.receive(), timeout=1) + + assert isinstance(replies, list), "both entries must travel back in one array frame" + assert [reply["result"] for reply in replies] == [ + {"n": 1, "peer": {"answered": 1}}, + {"n": 2, "peer": {"answered": 2}}, + ] + + +@pytest.mark.asyncio +async def test_cancelled_batch_entry_answers_once_and_keeps_its_sibling_result(): + started = asyncio.Event() + + async def handler(method: str, params: Any, is_notification: bool) -> Any: + if params == {"slow": True}: + started.set() + await asyncio.sleep(5) + return {"echo": params} + + async with _memory(handler) as (_, peer): + await peer.send([ + {"jsonrpc": "2.0", "id": "slow", "method": "probe", "params": {"slow": True}}, + {"jsonrpc": "2.0", "id": "fast", "method": "probe", "params": {"slow": False}}, + ]) + await asyncio.wait_for(started.wait(), timeout=1) + # The cancel names an id already registered from inside the same batch. + await peer.send([ + {"jsonrpc": "2.0", "method": PROTOCOL_METHODS["cancel_request"], "params": {"requestId": "slow"}} + ]) + + replies = await asyncio.wait_for(peer.receive(), timeout=1) + + assert [reply["id"] for reply in replies] == ["slow", "fast"] + assert (replies[0]["error"]["code"], replies[0]["error"]["message"]) == (-32800, "Request cancelled") + assert "result" not in replies[0] + assert replies[1] == {"jsonrpc": "2.0", "id": "fast", "result": {"echo": {"slow": False}}} + # A single array frame, so the cancelled id is answered exactly once. + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(peer.receive(), timeout=0.1) diff --git a/tests/test_protocol_negotiation.py b/tests/test_protocol_negotiation.py index 06e763e2..1a7937fb 100644 --- a/tests/test_protocol_negotiation.py +++ b/tests/test_protocol_negotiation.py @@ -235,5 +235,8 @@ def create_agent(connection: v2.AgentSideConnection) -> RoutedV2Agent: def _initialize_count(events: list[StreamEvent]) -> int: return sum( - event.direction == StreamDirection.OUTGOING and event.message.get("method") == "initialize" for event in events + event.direction == StreamDirection.OUTGOING + and isinstance(event.message, dict) + and event.message.get("method") == "initialize" + for event in events ) diff --git a/tests/test_request_cancellation.py b/tests/test_request_cancellation.py index 2206e7e8..a897bb70 100644 --- a/tests/test_request_cancellation.py +++ b/tests/test_request_cancellation.py @@ -106,14 +106,14 @@ class _GatedTransport: def __init__(self) -> None: self.incoming: asyncio.Queue[dict[str, Any]] = asyncio.Queue() - self.sent: list[dict[str, Any]] = [] + self.sent: list[dict[str, Any] | list[Any]] = [] self.sending = asyncio.Event() self.release = asyncio.Event() self.settled = asyncio.Event() self.send_cancelled = False self._receiving = False - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: self.sending.set() try: await self.release.wait() @@ -125,7 +125,7 @@ async def send(self, message: dict[str, Any]) -> None: finally: self.settled.set() - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> dict[str, Any] | list[Any] | None: self._receiving = True try: return await self.incoming.get() @@ -172,7 +172,7 @@ async def handler(method: str, params: Any, is_notification: bool) -> Any: assert not transport.send_cancelled, "the cancellation aborted the response send" if handler_cancelled: - assert [(m["id"], m["error"]["code"]) for m in transport.sent] == [(0, -32800)] + assert [_error_code_and_id(m) for m in transport.sent] == [(0, -32800)] else: assert transport.sent == [{"jsonrpc": "2.0", "id": 0, "result": {"ok": True}}] @@ -218,6 +218,11 @@ async def test_cancel_request_only_affects_the_targeted_request(server, caplog: await asyncio.wait_for(server.client_reader.readline(), timeout=0.1) +def _error_code_and_id(message: Any) -> tuple[Any, Any]: + assert isinstance(message, dict) + return message["id"], message["error"]["code"] + + async def _cancel_while_pending(request_id: Any, cancel: dict[str, Any]) -> dict[str, Any]: """Deliver ``cancel`` while request ``request_id`` is in flight, then let it finish and return its response.""" transport = _GatedTransport() @@ -238,6 +243,7 @@ async def handler(method: str, params: Any, is_notification: bool) -> Any: await asyncio.wait_for(transport.settled.wait(), timeout=1) [response] = transport.sent + assert isinstance(response, dict) return response diff --git a/tests/test_request_error_logging.py b/tests/test_request_error_logging.py index 12f3ec4d..1ac38bc3 100644 --- a/tests/test_request_error_logging.py +++ b/tests/test_request_error_logging.py @@ -9,7 +9,7 @@ from __future__ import annotations import logging -from typing import Any +from typing import Any, cast import pytest @@ -21,12 +21,12 @@ class _RecordingTransport: """Message transport that records outgoing frames.""" def __init__(self) -> None: - self.sent: list[dict[str, Any]] = [] + self.sent: list[dict[str, Any] | list[Any]] = [] - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: self.sent.append(message) - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> dict[str, Any] | list[Any] | None: return None async def close(self) -> None: @@ -70,7 +70,7 @@ async def test_run_request_unhandled_exception_is_logged_and_sent_as_internal_er # A handled application exception becomes exactly one JSON-RPC error frame. assert len(transport.sent) == 1 - response = transport.sent[0] + response = cast("dict[str, Any]", transport.sent[0]) assert response["id"] == 7 assert "result" not in response assert response["error"] == {"code": -32603, "message": "Internal error", "data": {"details": "kaboom"}} @@ -101,7 +101,7 @@ async def test_run_notification_unhandled_exception_is_logged_and_not_answered(c @pytest.mark.asyncio async def test_response_send_failure_is_not_mapped_to_a_handler_error(caplog): class _FailingTransport(_RecordingTransport): - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: raise ConnectionError("send failed") async def successful_handler(method: str, params: Any, is_notification: bool) -> Any: diff --git a/tests/test_rpc.py b/tests/test_rpc.py index 43c37fd1..1a68b08c 100644 --- a/tests/test_rpc.py +++ b/tests/test_rpc.py @@ -485,10 +485,10 @@ async def test_cancelled_request_is_removed_from_pending_map() -> None: sent = asyncio.Event() class _NeverRespondTransport: - async def send(self, message: dict[str, Any]) -> None: + async def send(self, message: dict[str, Any] | list[Any]) -> None: sent.set() - async def receive(self) -> dict[str, Any] | None: + async def receive(self) -> dict[str, Any] | list[Any] | None: await asyncio.Event().wait() return None diff --git a/tests/test_v2_runtime.py b/tests/test_v2_runtime.py index d93da0b8..80bbfe63 100644 --- a/tests/test_v2_runtime.py +++ b/tests/test_v2_runtime.py @@ -212,23 +212,31 @@ async def test_callable_agent_is_not_treated_as_a_factory() -> None: @pytest.mark.asyncio -async def test_v2_runtime_rejects_a_different_protocol_version() -> None: - client_transport, agent_transport = memory_transport_pair() +@pytest.mark.parametrize("requested_version", [1, v2.PROTOCOL_VERSION, 65535]) +async def test_v2_initialize_negotiates_the_supported_version_on_the_wire(requested_version: int) -> None: + # TCK ACP-INIT-003/201/202: a v2-only agent must answer 1, 2 and 65535 with version 2 instead + # of a -32602, and must still reject a second initialize on the same connection. + peer, agent_transport = memory_transport_pair() agent = Agent() agent_connection = v2.AgentSideConnection(agent, agent_transport) - client_connection = v2.ClientSideConnection(Client(), client_transport) + params = { + **initialize_request().model_dump(mode="json", by_alias=True, exclude_none=True), + "protocolVersion": requested_version, + } try: - with pytest.raises(RequestError) as error: - await client_connection.initialize( - protocol_version=1, info=v2.schema.Implementation(name="test-client", version="1.0.0") - ) + await peer.send({"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": params}) + response = await asyncio.wait_for(peer.receive(), timeout=1) - assert isinstance(error.value, RequestError) - assert error.value.code == -32602 - assert agent.initialize_calls == 0 + assert response["result"]["protocolVersion"] == v2.PROTOCOL_VERSION + assert agent.initialize_calls == 1 + + await peer.send({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": params}) + duplicate = await asyncio.wait_for(peer.receive(), timeout=1) + + assert (duplicate["id"], duplicate["error"]["code"]) == (1, -32600) + assert agent.initialize_calls == 1 finally: - await client_connection.close() await agent_connection.close() @@ -582,3 +590,49 @@ async def handle_extension_request(self, method: str, params: Any) -> Any: assert handled == [] finally: await agent_connection.close() + + +@pytest.mark.asyncio +async def test_v2_request_without_params_is_not_invalid_params() -> None: + # TCK ACP-BATCH-204/205/JSONRPC-001 batch ``session/list`` with no ``params`` member at all. + # Every ACP request envelope requires only ``id``/``method``, so an omitted ``params`` means + # "no arguments" and must not be rejected as Invalid params. + class ListAgent(Agent): + async def list_sessions( + self, cwd: str | None = None, cursor: str | None = None, **kwargs: Any + ) -> v2.schema.ListSessionsResponse: + return v2.schema.ListSessionsResponse(sessions=[]) + + async def new_session(self, cwd: str, **kwargs: Any) -> v2.schema.NewSessionResponse: + return v2.schema.NewSessionResponse(session_id="sess") + + peer, agent_transport = memory_transport_pair() + agent_connection = v2.AgentSideConnection(ListAgent(), agent_transport) + + try: + await peer.send({ + "jsonrpc": "2.0", + "id": 0, + "method": "initialize", + "params": initialize_request().model_dump(mode="json", by_alias=True, exclude_none=True), + }) + assert "result" in await asyncio.wait_for(peer.receive(), timeout=1) + + await peer.send({"jsonrpc": "2.0", "id": 1, "method": "session/list"}) + single = await asyncio.wait_for(peer.receive(), timeout=1) + assert single["result"] == {"sessions": []} + + await peer.send([ + {"jsonrpc": "2.0", "id": "a", "method": "session/list"}, + {"jsonrpc": "2.0", "id": "b", "method": "session/list"}, + ]) + batch = await asyncio.wait_for(peer.receive(), timeout=1) + assert [entry["id"] for entry in batch] == ["a", "b"] + assert all("result" in entry for entry in batch) + + # A method whose params really are required still fails validation. + await peer.send({"jsonrpc": "2.0", "id": 2, "method": "session/new"}) + rejected = await asyncio.wait_for(peer.receive(), timeout=1) + assert rejected["error"]["code"] == -32602 + finally: + await agent_connection.close()