Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
28 changes: 28 additions & 0 deletions src/acp/_jsonrpc.py
Original file line number Diff line number Diff line change
@@ -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()}
24 changes: 16 additions & 8 deletions src/acp/_sse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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:
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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
38 changes: 16 additions & 22 deletions src/acp/_transport.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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

Expand Down
4 changes: 3 additions & 1 deletion src/acp/http/asgi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"),
Expand Down
8 changes: 7 additions & 1 deletion src/acp/http/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down
10 changes: 7 additions & 3 deletions src/acp/http/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
22 changes: 8 additions & 14 deletions src/acp/ws/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
22 changes: 8 additions & 14 deletions src/acp/ws/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
32 changes: 31 additions & 1 deletion tests/http/test_asgi.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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:
Expand Down
Loading
Loading