Skip to content
Draft
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
11 changes: 9 additions & 2 deletions src/acp/_jsonrpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
15 changes: 8 additions & 7 deletions src/acp/_sse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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.
"""
Expand Down
24 changes: 12 additions & 12 deletions src/acp/_transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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: ...

Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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
62 changes: 60 additions & 2 deletions src/acp/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -33,14 +34,19 @@ 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]

_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))
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
10 changes: 5 additions & 5 deletions src/acp/experimental/v2/_initialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
6 changes: 5 additions & 1 deletion src/acp/experimental/v2/_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}
Expand Down
17 changes: 13 additions & 4 deletions src/acp/http/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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"))
Expand Down
2 changes: 1 addition & 1 deletion src/acp/task/sender.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
6 changes: 3 additions & 3 deletions src/acp/ws/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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()
Expand Down
6 changes: 3 additions & 3 deletions src/acp/ws/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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":
Expand Down
15 changes: 15 additions & 0 deletions tests/http/test_http_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading
Loading