fix: align JSON-RPC error handling across transports (#174)

Author: PsiACE

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()}

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

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
 

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

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

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)

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:

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:

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: