From 6925a38949a138285a12c5f8e6040eefc87bb8bf Mon Sep 17 00:00:00 2001 From: Chojan Shang Date: Sun, 18 Oct 2026 18:22:54 +0800 Subject: [PATCH] fix: align JSON-RPC error handling across transports --- src/acp/_jsonrpc.py | 28 ++++++ src/acp/_sse.py | 24 +++-- src/acp/_transport.py | 38 ++++---- src/acp/http/asgi.py | 4 +- src/acp/http/client.py | 8 +- src/acp/http/server.py | 10 ++- src/acp/ws/client.py | 22 ++--- src/acp/ws/server.py | 22 ++--- tests/http/test_asgi.py | 32 ++++++- tests/http/test_http_client.py | 26 ++++++ tests/http/test_http_server.py | 4 +- tests/http/test_transport_errors.py | 133 ++++++++++++++++++++++++++++ 12 files changed, 286 insertions(+), 65 deletions(-) create mode 100644 src/acp/_jsonrpc.py create mode 100644 tests/http/test_transport_errors.py diff --git a/src/acp/_jsonrpc.py b/src/acp/_jsonrpc.py new file mode 100644 index 00000000..078c5621 --- /dev/null +++ b/src/acp/_jsonrpc.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import json +from typing import Any + +from .exceptions import RequestError + + +def decode_message(payload: str | bytes) -> dict[str, Any] | None: + """Decode a message, returning None for unsupported non-empty batches.""" + try: + message = json.loads(payload) + except (json.JSONDecodeError, UnicodeDecodeError) as exc: + raise RequestError.parse_error() from exc + return require_message_object(message) + + +def require_message_object(message: Any) -> dict[str, Any] | None: + """Reject invalid root values while leaving non-empty batches unsupported.""" + if isinstance(message, dict) and message: + return message + if isinstance(message, list) and message: + return None + raise RequestError.invalid_request() + + +def error_response(error: RequestError) -> dict[str, Any]: + return {"jsonrpc": "2.0", "id": None, "error": error.to_error_obj()} diff --git a/src/acp/_sse.py b/src/acp/_sse.py index 73f868c9..2f67e3e2 100644 --- a/src/acp/_sse.py +++ b/src/acp/_sse.py @@ -11,8 +11,11 @@ import json from typing import TYPE_CHECKING, Any +from ._jsonrpc import decode_message +from .exceptions import RequestError + if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncIterator, Awaitable, Callable __all__ = [ "parse_sse_stream", @@ -36,16 +39,19 @@ def serialize_sse_keepalive() -> bytes: return b": keepalive\n\n" -def _decode_event(data_lines: list[str]) -> dict[str, Any] | None: +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.""" payload = "\n".join(data_lines) if not payload: return None try: - message = json.loads(payload) - except json.JSONDecodeError: + return decode_message(payload) + except RequestError as error: + if on_error is not None: + await on_error(error) return None - return message if isinstance(message, dict) else None def _append_field(line: str, data_lines: list[str]) -> None: @@ -57,7 +63,9 @@ def _append_field(line: str, data_lines: list[str]) -> None: data_lines.append(value) -async def parse_sse_stream(chunks: AsyncIterator[bytes]) -> AsyncIterator[dict[str, Any]]: +async def parse_sse_stream( + chunks: AsyncIterator[bytes], *, on_error: Callable[[RequestError], Awaitable[None]] | None = None +) -> AsyncIterator[dict[str, Any]]: """Parse an SSE byte stream, yielding decoded JSON-RPC ``data:`` payloads. Comments (lines starting with ``:``) and non-``data`` fields are ignored. @@ -75,7 +83,7 @@ async def parse_sse_stream(chunks: AsyncIterator[bytes]) -> AsyncIterator[dict[s line, buffer = buffer.split("\n", 1) line = line.rstrip("\r") if line == "": - event = _decode_event(data_lines) + event = await _decode_event(data_lines, on_error) data_lines = [] if event is not None: yield event @@ -85,6 +93,6 @@ async def parse_sse_stream(chunks: AsyncIterator[bytes]) -> AsyncIterator[dict[s decoder.decode(b"", final=True) # Flush a trailing event with no terminating blank line. - event = _decode_event(data_lines) + event = await _decode_event(data_lines, on_error) if event is not None: yield event diff --git a/src/acp/_transport.py b/src/acp/_transport.py index 1ded5e97..2e2d113a 100644 --- a/src/acp/_transport.py +++ b/src/acp/_transport.py @@ -1,24 +1,19 @@ """Message-level transport seam for :class:`acp.connection.Connection`. -Historically the connection spoke directly to ``asyncio`` byte streams with -newline framing. To support message-oriented remote transports (Streamable -HTTP + WebSocket) we introduce a small :class:`Transport` protocol that moves -JSON-RPC *messages* (already-decoded ``dict`` payloads) instead of bytes. - -The existing stdio path is re-expressed on top of this seam via -:class:`NdjsonTransport`, which wraps the current byte-stream framing so there -is **zero behaviour change** for stdio users. :func:`memory_transport_pair` -gives two linked in-memory transports for in-process connections and tests. +Transports exchange decoded JSON-RPC objects. :class:`NdjsonTransport` provides +newline framing for byte streams, while :func:`memory_transport_pair` connects +two peers in memory. """ from __future__ import annotations import asyncio import contextlib -import json -import logging from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from ._jsonrpc import decode_message, error_response +from .exceptions import RequestError + if TYPE_CHECKING: from collections.abc import Awaitable @@ -49,12 +44,7 @@ async def close(self) -> None: ... class NdjsonTransport: - """Transport backed by newline-delimited JSON over asyncio byte streams. - - This preserves the exact framing, buffer-limit-overrun handling, and receive - timeout semantics that previously lived inside ``Connection`` so that stdio - behaviour is byte-for-byte unchanged. - """ + """Transport backed by newline-delimited JSON over asyncio byte streams.""" def __init__( self, @@ -66,6 +56,7 @@ def __init__( self._reader = reader self._sender = sender self._receive_timeout = receive_timeout + self._error_send_failed = False async def send(self, message: dict[str, Any]) -> None: await self._sender.send(message) @@ -79,12 +70,15 @@ async def receive(self) -> dict[str, Any] | None: if not line: continue try: - message = json.loads(line) - except Exception: - logging.exception("Error parsing JSON-RPC message") + message = decode_message(line) + except RequestError as error: + if not self._error_send_failed: + try: + await self.send(error_response(error)) + except ConnectionError: + self._error_send_failed = True continue - if not isinstance(message, dict): - logging.warning("Ignoring non-object JSON-RPC message") + if message is None: continue return message diff --git a/src/acp/http/asgi.py b/src/acp/http/asgi.py index 5e8216cc..4f696821 100644 --- a/src/acp/http/asgi.py +++ b/src/acp/http/asgi.py @@ -18,6 +18,8 @@ from starlette.routing import Route, WebSocketRoute from starlette.websockets import WebSocket +from .._jsonrpc import error_response +from ..exceptions import RequestError from ..ws.server import handle_websocket from .protocol import ACP_ENDPOINT_PATH, CONNECTION_ID_HEADER, CONTENT_TYPE_SSE, SESSION_ID_HEADER from .server import AcpServer, AgentFactory @@ -62,7 +64,7 @@ async def _post(server: AcpServer, request: Request) -> Response: try: message = await request.json() except (json.JSONDecodeError, UnicodeDecodeError): - return JSONResponse({"error": "Invalid JSON"}, status_code=400) + return JSONResponse(error_response(RequestError.parse_error()), status_code=400) return await server.handle_post( message, content_type=request.headers.get("content-type"), diff --git a/src/acp/http/client.py b/src/acp/http/client.py index 8baf6bdf..5087bcf9 100644 --- a/src/acp/http/client.py +++ b/src/acp/http/client.py @@ -23,7 +23,9 @@ import contextlib from typing import TYPE_CHECKING, Any +from .._jsonrpc import error_response from .._sse import parse_sse_stream +from ..exceptions import RequestError from .protocol import ( CONNECTION_ID_HEADER, CONTENT_TYPE_JSON, @@ -190,7 +192,7 @@ async def _consume_stream(self, *, session_id: str | None) -> None: async with self._client.stream("GET", self._url, headers=headers) as response: if response.status_code != 200: return - async for event in parse_sse_stream(_aiter_raw(response)): + async for event in parse_sse_stream(_aiter_raw(response), on_error=self._send_protocol_error): self._handle_incoming(event) except (httpx2.HTTPError, asyncio.CancelledError): return @@ -215,6 +217,10 @@ def _on_stream_closed(self, session_id: str | None) -> None: self._pending_loads.clear() self._inbox.put_nowait(_EOF) + 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: # Load and resume responses may omit sessionId; the ID is in the request. if is_response_message(message): diff --git a/src/acp/http/server.py b/src/acp/http/server.py index 049d0372..7d88d95c 100644 --- a/src/acp/http/server.py +++ b/src/acp/http/server.py @@ -29,8 +29,10 @@ except ImportError as exc: raise ImportError("The HTTP server requires the 'http' extra: pip install agent-client-protocol[http]") from exc +from .._jsonrpc import error_response, require_message_object from .._sse import serialize_sse_event, serialize_sse_keepalive from ..agent.connection import AgentSideConnection +from ..exceptions import RequestError from .protocol import ( CONNECTION_ID_HEADER, SESSION_ATTACH_METHODS, @@ -249,10 +251,12 @@ async def handle_post( ) -> Response: if content_type is None or not content_type.lower().startswith("application/json"): return JSONResponse({"error": "Content-Type must be application/json"}, status_code=415) - if isinstance(message, list): + try: + message = require_message_object(message) + except RequestError as error: + return JSONResponse(error_response(error), status_code=400) + if message is None: return JSONResponse({"error": "Batch requests are not supported"}, status_code=501) - if not isinstance(message, dict): - return JSONResponse({"error": "Invalid JSON-RPC message"}, status_code=400) if is_initialize_request(message): return await self._handle_initialize(message) diff --git a/src/acp/ws/client.py b/src/acp/ws/client.py index 531739f7..2ac2da7a 100644 --- a/src/acp/ws/client.py +++ b/src/acp/ws/client.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any from .._cookies import MemoryAcpCookieStore +from .._jsonrpc import decode_message, error_response from ..exceptions import RequestError try: @@ -54,21 +55,14 @@ async def receive(self) -> dict[str, Any] | None: if isinstance(frame, bytes): continue try: - message = json.loads(frame) - except json.JSONDecodeError: - error = RequestError.parse_error() - else: - if isinstance(message, dict) and message: + message = decode_message(frame) + if message is not None: return message - if isinstance(message, list) and message: - # Batch support is outside this message-level transport. - continue - error = RequestError.invalid_request() - # Sending while closing waits for EOF, which may require draining - # buffered frames first. Keep receiving instead of waiting here. - if self._ws.state == websockets.State.OPEN: - with contextlib.suppress(websockets.ConnectionClosed): - await self.send({"jsonrpc": "2.0", "id": None, "error": error.to_error_obj()}) + except RequestError as error: + # Sending while closing may wait for buffered frames to drain. + if self._ws.state == websockets.State.OPEN: + with contextlib.suppress(websockets.ConnectionClosed): + await self.send(error_response(error)) async def close(self) -> None: if self._closed: diff --git a/src/acp/ws/server.py b/src/acp/ws/server.py index 543a790d..0ae5dcb2 100644 --- a/src/acp/ws/server.py +++ b/src/acp/ws/server.py @@ -3,13 +3,13 @@ from __future__ import annotations import contextlib -import json import logging import uuid from typing import TYPE_CHECKING, Any from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState +from .._jsonrpc import decode_message, error_response from ..agent.connection import AgentSideConnection from ..exceptions import RequestError from ..http.protocol import CONNECTION_ID_HEADER @@ -44,20 +44,14 @@ async def receive(self) -> dict[str, Any] | None: logging.warning("Ignoring binary WebSocket frame; ACP messages must be sent as text frames") continue try: - message = json.loads(event["text"]) - except json.JSONDecodeError: - error = RequestError.parse_error() - else: - if isinstance(message, dict) and message: + message = decode_message(event["text"]) + if message is not None: return message - if isinstance(message, list) and message: - # Batch support is outside this message-level transport. - continue - error = RequestError.invalid_request() - # A failed write must not discard valid messages buffered by ASGI. - if self._ws.application_state == WebSocketState.CONNECTED: - with contextlib.suppress(WebSocketDisconnect): - await self.send({"jsonrpc": "2.0", "id": None, "error": error.to_error_obj()}) + except RequestError as error: + # A failed write must not discard valid messages buffered by ASGI. + if self._ws.application_state == WebSocketState.CONNECTED: + with contextlib.suppress(WebSocketDisconnect): + await self.send(error_response(error)) return None async def close(self) -> None: diff --git a/tests/http/test_asgi.py b/tests/http/test_asgi.py index b9ee3e40..97ccb7d3 100644 --- a/tests/http/test_asgi.py +++ b/tests/http/test_asgi.py @@ -28,7 +28,8 @@ ("POST", {"Content-Type": "text/plain"}, "{}", 415), ("POST", {"Content-Type": "application/json"}, "invalid", 400), ("POST", {"Content-Type": "application/json"}, "null", 400), - ("POST", {"Content-Type": "application/json"}, "[]", 501), + ("POST", {"Content-Type": "application/json"}, "[]", 400), + ("POST", {"Content-Type": "application/json"}, '[{"jsonrpc":"2.0","method":"probe","id":1}]', 501), ("GET", {"Accept": "application/json"}, "", 406), ("GET", {"Accept": "text/event-stream"}, "", 400), ("GET", {"Accept": "text/event-stream", CONNECTION_ID_HEADER: "unknown"}, "", 404), @@ -47,6 +48,35 @@ async def test_http_errors(method: str, headers: dict[str, str], body: str, stat assert "error" in response.json() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "code"), + [ + ("invalid json", -32700), + ("null", -32600), + ("[]", -32600), + ("{}", -32600), + ("true", -32600), + ("42", -32600), + ('"text"', -32600), + ], +) +async def test_invalid_post_returns_protocol_error_without_losing_connection(payload: str, code: int) -> None: + app = create_asgi_app(lambda conn: TestAgent()) + async with ( + app.router.lifespan_context(app), + httpx2.AsyncClient(transport=httpx2.ASGITransport(app=app), base_url="http://test") as client, + ): + initialized = await client.post("/acp", json=INITIALIZE) + headers = {"Content-Type": "application/json", CONNECTION_ID_HEADER: initialized.headers[CONNECTION_ID_HEADER]} + response = await client.post("/acp", content=payload, headers=headers) + assert response.status_code == 400 + body = response.json() + assert body["jsonrpc"] == "2.0" and body["id"] is None + assert body["error"]["code"] == code + assert (await client.delete("/acp", headers=headers)).status_code == 202 + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["PUT", "OPTIONS"]) async def test_unsupported_methods_do_not_open_a_stream(method: str) -> None: diff --git a/tests/http/test_http_client.py b/tests/http/test_http_client.py index daaa1f6a..38ddb78b 100644 --- a/tests/http/test_http_client.py +++ b/tests/http/test_http_client.py @@ -154,6 +154,32 @@ async def test_connection_scoped_sse_skips_non_object_json(payload: bytes) -> No await client.aclose() +@pytest.mark.asyncio +async def test_rejected_protocol_error_does_not_discard_buffered_response() -> None: + server = FakeServer() + + def handle(request: httpx2.Request) -> httpx2.Response: + response = server.handler(request) + if request.method == "POST" and "error" in json.loads(request.content): + return httpx2.Response(503) + return response + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: + transport = create_http_stream("http://test/acp", client=client) + try: + await transport.send({"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {}}) + await transport.receive() + server.conn_stream.put_nowait(b"data: invalid json\n\n") + response = {"jsonrpc": "2.0", "id": 1, "result": {"items": [], "value": None}} + server.push_conn(response) + server.conn_stream.put_nowait(None) + assert await asyncio.wait_for(transport.receive(), timeout=1) == response + assert await asyncio.wait_for(transport.receive(), timeout=1) is None + assert server.posts[-1]["error"]["code"] == -32700 + finally: + await transport.close() + + @pytest.mark.asyncio async def test_session_scoped_sse_opens_after_new_session() -> None: server = FakeServer() diff --git a/tests/http/test_http_server.py b/tests/http/test_http_server.py index 382579e1..6ab739b5 100644 --- a/tests/http/test_http_server.py +++ b/tests/http/test_http_server.py @@ -71,7 +71,9 @@ async def test_post_wrong_content_type_returns_415() -> None: @pytest.mark.asyncio async def test_batch_returns_501() -> None: server = AcpServer(_agent_factory(_Agent())) - result = await server.handle_post([], content_type=CT_JSON, connection_id=None, session_id=None) + result = await server.handle_post( + [{"jsonrpc": "2.0", "method": "probe", "id": 1}], content_type=CT_JSON, connection_id=None, session_id=None + ) assert result.status_code == 501 await server.close() diff --git a/tests/http/test_transport_errors.py b/tests/http/test_transport_errors.py new file mode 100644 index 00000000..da2d47ed --- /dev/null +++ b/tests/http/test_transport_errors.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass +from typing import Any + +import httpx2 +import pytest +import pytest_asyncio +from websockets.asyncio.server import serve + +from acp.connection import Connection +from tests.http.test_http_client import FakeServer + + +@dataclass +class _Peer: + connection: Connection + send: Callable[[str], Awaitable[None]] + receive: Callable[[], Awaitable[dict[str, Any]]] + + +async def _handler(method: str, params: Any, is_notification: bool) -> None: + pass + + +@pytest_asyncio.fixture(params=["stdio", "websocket", "sse"]) +async def peer(request: pytest.FixtureRequest, server) -> AsyncIterator[_Peer]: + if request.param == "stdio": + + async def send(text: str) -> None: + server.server_writer.write((text + "\n").encode()) + await server.server_writer.drain() + + async def receive() -> dict[str, Any]: + return json.loads(await server.server_reader.readline()) + + async with Connection(_handler, server.client_writer, server.client_reader) as connection: + yield _Peer(connection, send, receive) + elif request.param == "websocket": + from acp.ws.client import create_websocket_stream + + socket_ready = asyncio.get_running_loop().create_future() + + async def accept(socket: Any) -> None: + socket_ready.set_result(socket) + await socket.wait_closed() + + async with serve(accept, "127.0.0.1", 0) as listener: + transport = await create_websocket_stream(f"ws://127.0.0.1:{listener.sockets[0].getsockname()[1]}") + socket = await socket_ready + + async def receive() -> dict[str, Any]: + return json.loads(await socket.recv()) + + async with Connection(_handler, transport) as connection: + yield _Peer(connection, socket.send, receive) + else: + fake = FakeServer() + posts: asyncio.Queue[dict[str, Any]] = asyncio.Queue() + + async def handle(request: httpx2.Request) -> httpx2.Response: + if request.method == "POST": + posts.put_nowait(json.loads(request.content)) + return fake.handler(request) + + from acp.http.client import create_http_stream + + async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client: + transport = create_http_stream("http://test/acp", client=client) + await transport.send({"jsonrpc": "2.0", "id": "initialize", "method": "initialize", "params": {}}) + await transport.receive() + await posts.get() + + async def send(text: str) -> None: + fake.conn_stream.put_nowait(("data: " + text + "\n\n").encode()) + + async with Connection(_handler, transport) as connection: + yield _Peer(connection, send, posts.get) + + +async def _request(peer: _Peer) -> tuple[asyncio.Task[Any], dict[str, Any]]: + pending = asyncio.create_task(peer.connection.send_request("probe", {})) + request = await asyncio.wait_for(peer.receive(), timeout=1) + assert request["method"] == "probe" + return pending, request + + +async def _complete(peer: _Peer, pending: asyncio.Task[Any], request: dict[str, Any]) -> None: + values = [None, [], {}, False, 0, "", {"nested": ["雪🙂"]}] + await peer.send(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": values})) + assert await asyncio.wait_for(pending, timeout=1) == values + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "code"), + [ + ("invalid json", -32700), + ("null", -32600), + ("[]", -32600), + ("{}", -32600), + ("42", -32600), + ("1.5", -32600), + ('"text"', -32600), + ("true", -32600), + ("false", -32600), + ], +) +async def test_invalid_payload_reports_error_and_preserves_requests(peer: _Peer, payload: str, code: int) -> None: + pending, request = await _request(peer) + await peer.send(payload) + error = await asyncio.wait_for(peer.receive(), timeout=1) + assert error["jsonrpc"] == "2.0" and error["id"] is None + assert error["error"]["code"] == code + assert not pending.done() + await _complete(peer, pending, request) + next_pending, next_request = await _request(peer) + await _complete(peer, next_pending, next_request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", ['{"jsonrpc":"2.0","id":[],"result":null}', '[{"jsonrpc":"2.0","id":123,"result":null}]'] +) +async def test_ignored_payload_preserves_requests_without_reply(peer: _Peer, payload: str) -> None: + pending, request = await _request(peer) + await peer.send(payload) + await _complete(peer, pending, request) + next_pending, next_request = await _request(peer) + await _complete(peer, next_pending, next_request)