From 5ae0734513979bc2cefed6fa0da7ae58723b4beb Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Fri, 11 Sep 2026 11:04:25 +0800 Subject: [PATCH 1/2] refactor: simplify web transports with Starlette --- docs/web-transport.md | 115 ++++++++++- examples/http_server.py | 2 +- pyproject.toml | 3 +- src/acp/_transport.py | 7 +- src/acp/http/__init__.py | 2 +- src/acp/http/asgi.py | 252 ++++++++--------------- src/acp/http/protocol.py | 2 +- src/acp/http/server.py | 362 +++++++++++---------------------- src/acp/ws/server.py | 116 +++++------ tests/http/test_asgi.py | 178 ++++++++++++++++ tests/http/test_fixes.py | 59 +++++- tests/http/test_http_server.py | 98 +++++++-- tests/http/test_websocket.py | 75 ++++++- uv.lock | 17 ++ 14 files changed, 760 insertions(+), 528 deletions(-) create mode 100644 tests/http/test_asgi.py diff --git a/docs/web-transport.md b/docs/web-transport.md index e601988..a6fb166 100644 --- a/docs/web-transport.md +++ b/docs/web-transport.md @@ -21,7 +21,9 @@ Both reuse the existing JSON-RPC message format and ACP lifecycle pip install "agent-client-protocol[http]" ``` -This pulls in `httpx[http2]` (HTTP/2 + SSE consumption) and `websockets`. +This pulls in `httpx[http2]` (HTTP/2 + SSE consumption), `websockets`, and +`starlette` (the server application). The core SDK and stdio transport do not +require these optional dependencies. ## Client @@ -55,8 +57,8 @@ stream; reconnect/retry is the caller's responsibility (v1 of the RFD). ## Server -The server core is framework-agnostic; a thin ASGI adapter bridges it to your -web framework: +The server uses Starlette for HTTP requests, responses, routing, streaming, +WebSocket handling, and application lifespan: ```python from acp.http.asgi import create_asgi_app @@ -65,8 +67,111 @@ from acp.http.asgi import create_asgi_app app = create_asgi_app(lambda conn: MyAgent()) ``` -`app` is a standard ASGI 3.0 application handling `POST`/`GET`/`DELETE` and -WebSocket upgrades on the ACP endpoint. +`app` is a `starlette.applications.Starlette` instance handling +`POST`/`GET`/`DELETE` and WebSocket upgrades at `/acp` by default. Set the +keyword-only `path` argument to use a different endpoint for both transports: + +```python +app = create_asgi_app(lambda conn: MyAgent(), path="/rpc") +``` + +Other paths do not serve ACP. Starlette supplies +`Request`, `JSONResponse`, `StreamingResponse`, and `WebSocket`; the SDK keeps +ACP connection and session routing. + +### Mounting in another application + +Mount the app at the desired prefix. The parent must enter the child lifespan +so HTTP connections are cleaned up during shutdown (mounted application +lifespans are not run automatically): + +```python +from contextlib import asynccontextmanager +from starlette.applications import Starlette +from starlette.routing import Mount + +acp_app = create_asgi_app(lambda conn: MyAgent()) + +@asynccontextmanager +async def lifespan(app): + async with acp_app.router.lifespan_context(acp_app): + yield + +app = Starlette(routes=[Mount("/agents", app=acp_app)], lifespan=lifespan) +# Connect to /agents/acp using either HTTP or WebSocket. +``` + +With `path="/rpc"`, the mounted endpoint is `/agents/rpc`. Use `path="/"` to +serve ACP at the mount root (`/agents/`). + +### How the server fits together + +Start reading at `acp/http/asgi.py`. It creates Starlette routes, passes parsed +HTTP requests to `AcpServer`, and binds WebSockets in `acp/ws/server.py`. Both use the existing +`AgentSideConnection` and its message-level `Transport` interface: + +```text +HTTP POST → _HttpTransport incoming queue → AgentSideConnection → agent +HTTP GET ← StreamingResponse ← SSE buffer ← _HttpTransport.send() ← agent output + +Starlette WebSocket ↔ _WebSocketTransport ↔ AgentSideConnection ↔ agent +``` + +For HTTP, `AcpServer` owns a dictionary of active connections. Each connection +has one incoming queue and one SSE buffer per stream. The incoming queue lets +POST return `202` while the agent handles the request. Output goes directly to +the relevant SSE buffer; there is no intermediate transport pair or pump task. + +HTTP output needs three routing rules: + +| Message | Destination | Why | +| --- | --- | --- | +| `initialize` response | POST body, via one Future | Establishes the connection before GET streams open | +| Response containing a new `sessionId` | Connection SSE stream | The client needs the ID before it can open the session stream | +| Other messages | Session SSE stream when known, otherwise connection stream | Responses use their request's recorded session; requests/notifications carry `sessionId` | + +`OutboundStream` retains a bounded buffer, backpressure, and close handling. +Idle SSE streams emit keepalives. These support slow readers, streams that open +after messages arrive, and orderly teardown. `DELETE` and server shutdown close +the HTTP connections and cancel their agent work. + +WebSocket already provides one bidirectional stream. Its transport adapts +Starlette's socket to JSON-RPC messages; it needs no HTTP connection registry, +session routing, SSE buffers, or multiplex mode. The ASGI handler owns the agent +connection and closes it on socket disconnect or handler cancellation. + +### Simplification experiment + +The original 80 HTTP, WebSocket, and RPC tests passed after each ablation. +The Starlette migration also passes these behaviors; assertions now inspect +Starlette response objects and WebSocket tests use the framework's socket: + +| Stage | Removed | Lines across the three server files | +| --- | --- | ---: | +| Baseline | — | 715 | +| First ablation | `ConnectionRegistry`, WebSocket multiplex mode, WebSocket pump tasks, forwarding-only ASGI method | 647 | +| Second ablation | HTTP memory transport pair and pump, `ConnectionState`, generic response-waiter map | 581 | +| Starlette migration | Custom ASGI app, request/header parsing, response encoding, WebSocket state tracking, `PostResult` | 483 | + +This measures structural simplification and regression coverage, not throughput +or latency. Additional tests cover interrupted initialization, closing a full +SSE buffer, WebSocket cancellation/disconnect, invalid frames, and session routing +of concurrent success/error responses. + +`create_asgi_app(agent_factory, *, path="/acp")` returns a Starlette application. +The default route is now `/acp`, replacing the earlier catch-all route. +`AcpServer.handle_post()` and `handle_delete()` return +Starlette responses (`status_code`, byte `body`, and case-insensitive `headers`). +`open_stream()` and `close()` keep their signatures. +The experimental `AcpAsgiApp` and `PostResult` wrappers were removed, along with +`ConnectionRegistry`, `ConnectionState`, `AcpServer.registry`, and +`create_websocket_connection()`. Direct WebSocket integrations now use +`handle_websocket(agent_factory, websocket)` with a Starlette `WebSocket`. +WebSocket lifetimes belong to their ASGI handlers; `AcpServer.close()` manages +HTTP connections. + +The migration adds checks for HTTP error statuses, unsupported methods, mounting, +lifespan cleanup, and reopening an SSE stream after disconnect. ### HTTP/2 server requirement diff --git a/examples/http_server.py b/examples/http_server.py index 624bcee..306de25 100644 --- a/examples/http_server.py +++ b/examples/http_server.py @@ -59,7 +59,7 @@ async def prompt(self, session_id: str, prompt: list[Any], **kwargs: Any) -> Pro return PromptResponse(stop_reason="end_turn") -# One agent instance per connection. +# A Starlette application with one agent instance per connection. app = create_asgi_app(lambda conn: EchoAgent()) diff --git a/pyproject.toml b/pyproject.toml index ffbd061..5722c3a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,12 +48,13 @@ dev = [ "httpx[http2]>=0.27", "websockets>=12.0", "uvicorn>=0.30", + "starlette>=0.49.3", ] [project.optional-dependencies] logfire = ["logfire>=0.14", "opentelemetry-sdk>=1.28.0"] # Experimental remote transports (Streamable HTTP + WebSocket), client + server. -http = ["httpx[http2]>=0.27", "websockets>=12.0"] +http = ["httpx[http2]>=0.27", "websockets>=12.0", "starlette>=0.49.3"] [build-system] requires = ["pdm-backend"] diff --git a/src/acp/_transport.py b/src/acp/_transport.py index 17c2815..9aa7ecc 100644 --- a/src/acp/_transport.py +++ b/src/acp/_transport.py @@ -8,8 +8,7 @@ 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, used by the HTTP/WS server to bind an -``AgentSideConnection`` to its message pump. +gives two linked in-memory transports for in-process connections and tests. """ from __future__ import annotations @@ -140,9 +139,7 @@ def memory_transport_pair() -> tuple[Transport, Transport]: """Return two linked in-memory transports. A message ``send`` on one end becomes available via ``receive`` on the - other. Closing an end enqueues an EOF (``None``) for its peer. This mirrors - the ``TransformStream`` pair the TypeScript SDK uses to bind a server-side - connection to its HTTP/WS message pump. + 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() diff --git a/src/acp/http/__init__.py b/src/acp/http/__init__.py index 7529047..dedcbd3 100644 --- a/src/acp/http/__init__.py +++ b/src/acp/http/__init__.py @@ -1,7 +1,7 @@ """Streamable HTTP transport for ACP (experimental). Public exports are import-guarded: the heavy client/server implementations pull -in optional dependencies (``httpx[http2]``). Importing a symbol without the +in optional dependencies (``httpx[http2]`` and ``starlette``). Importing a symbol without the extra installed raises a friendly ``ImportError`` pointing at ``pip install agent-client-protocol[http]``. """ diff --git a/src/acp/http/asgi.py b/src/acp/http/asgi.py index 7dbd5f8..84e59a3 100644 --- a/src/acp/http/asgi.py +++ b/src/acp/http/asgi.py @@ -1,175 +1,91 @@ -"""Thin ASGI adapter bridging Starlette/FastAPI/Hypercorn to :class:`AcpServer`. +"""Starlette application for ACP over Streamable HTTP and WebSocket. -``create_asgi_app(agent_factory)`` returns an ASGI 3.0 application callable that -handles POST/GET/DELETE (and WebSocket upgrades) on the ACP endpoint. Users can -mount it directly or wrap it in their framework of choice. - -Note: for a spec-compliant Streamable HTTP server, run this under an -HTTP/2-capable ASGI server (Hypercorn, Daphne, Granian) or terminate HTTP/2 at a -proxy. Uvicorn does not serve HTTP/2 (WebSocket still works). +Run the application directly or mount it under another ASGI application. +Use an HTTP/2-capable ASGI server for Streamable HTTP. """ from __future__ import annotations import json -from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -from .protocol import CONNECTION_ID_HEADER, CONTENT_TYPE_SSE, SESSION_ID_HEADER -from .server import AcpServer - -if TYPE_CHECKING: - from .server import AgentFactory - -__all__ = ["AcpAsgiApp", "create_asgi_app"] - -_JSON_HEADERS = [(b"content-type", b"application/json")] - - -def _header_lookup(scope_headers: list[tuple[bytes, bytes]], name: str) -> str | None: - target = name.lower().encode() - for key, value in scope_headers: - if key.lower() == target: - return value.decode("latin-1") - return None - - -class AcpAsgiApp: - """ASGI application wrapping an :class:`AcpServer`.""" - - def __init__(self, server: AcpServer) -> None: - self._server = server - - async def __call__(self, scope: dict[str, Any], receive: Callable, send: Callable) -> None: - scope_type = scope["type"] - if scope_type == "lifespan": - await self._handle_lifespan(receive, send) - return - if scope_type == "websocket": - await self._handle_websocket(scope, receive, send) - return - if scope_type != "http": - return - - method = scope["method"] - if method == "POST": - await self._handle_post(scope, receive, send) - elif method == "GET": - await self._handle_get(scope, receive, send) - elif method == "DELETE": - await self._handle_delete(scope, send) - else: - await self._send_json(send, 405, {"error": "Method not allowed"}) - - async def _handle_lifespan(self, receive: Callable, send: Callable) -> None: - while True: - message = await receive() - if message["type"] == "lifespan.startup": - await send({"type": "lifespan.startup.complete"}) - elif message["type"] == "lifespan.shutdown": - await self._server.close() - await send({"type": "lifespan.shutdown.complete"}) - return - - async def _read_body(self, receive: Callable) -> bytes: - chunks: list[bytes] = [] - while True: - message = await receive() - if message["type"] == "http.request": - chunks.append(message.get("body", b"")) - if not message.get("more_body", False): - break - elif message["type"] == "http.disconnect": - break - return b"".join(chunks) - - async def _handle_post(self, scope: dict[str, Any], receive: Callable, send: Callable) -> None: - headers = scope["headers"] - content_type = _header_lookup(headers, "content-type") - connection_id = _header_lookup(headers, CONNECTION_ID_HEADER) - session_id = _header_lookup(headers, SESSION_ID_HEADER) - raw = await self._read_body(receive) - try: - message = json.loads(raw) if raw else None - except json.JSONDecodeError: - await self._send_json(send, 400, {"error": "Invalid JSON"}) - return - result = await self._server.handle_post( - message, - content_type=content_type, - connection_id=connection_id, - session_id=session_id, - ) - await self._send_json(send, result.status, result.body, extra_headers=result.headers) - - async def _handle_get(self, scope: dict[str, Any], receive: Callable, send: Callable) -> None: - headers = scope["headers"] - upgrade = _header_lookup(headers, "upgrade") - if upgrade is not None and upgrade.lower() == "websocket": - # WebSocket upgrades arrive as scope type "websocket" in ASGI; a GET - # http scope with Upgrade is non-standard, so reject clearly. - await self._send_json(send, 400, {"error": "WebSocket upgrade must use the ws scope"}) - return - accept = _header_lookup(headers, "accept") or "" - if CONTENT_TYPE_SSE not in accept and "*/*" not in accept: - await self._send_json(send, 406, {"error": "Accept must include text/event-stream"}) - return - connection_id = _header_lookup(headers, CONNECTION_ID_HEADER) - session_id = _header_lookup(headers, SESSION_ID_HEADER) - error = self._server.validate_stream(connection_id=connection_id, session_id=session_id) - if error is not None: - await self._send_json(send, error.status, error.body) - return - if connection_id is None: # validated above, narrow for type-checker - await self._send_json(send, 400, {"error": "Missing connection id"}) - return - await send({ - "type": "http.response.start", - "status": 200, - "headers": [ - (b"content-type", CONTENT_TYPE_SSE.encode()), - (b"cache-control", b"no-cache"), - (b"connection", b"keep-alive"), - ], - }) - async for frame in self._server.open_stream(connection_id=connection_id, session_id=session_id): - await send({"type": "http.response.body", "body": frame, "more_body": True}) - await send({"type": "http.response.body", "body": b"", "more_body": False}) - - async def _handle_delete(self, scope: dict[str, Any], send: Callable) -> None: - connection_id = _header_lookup(scope["headers"], CONNECTION_ID_HEADER) - result = await self._server.handle_delete(connection_id=connection_id) - await self._send_json(send, result.status, result.body, extra_headers=result.headers) - - async def _handle_websocket(self, scope: dict[str, Any], receive: Callable, send: Callable) -> None: - from ..ws.server import handle_asgi_websocket - - await handle_asgi_websocket(self._server, scope, receive, send) - - async def _send_json( - self, - send: Callable, - status: int, - body: dict[str, Any] | None, - *, - extra_headers: dict[str, str] | None = None, - ) -> None: - payload = json.dumps(body).encode() if body is not None else b"" - headers = list(_JSON_HEADERS) - if extra_headers: - headers.extend((k.encode("latin-1"), v.encode("latin-1")) for k, v in extra_headers.items()) - await send({"type": "http.response.start", "status": status, "headers": headers}) - await send({"type": "http.response.body", "body": payload}) - - -def create_asgi_app(agent_factory: AgentFactory) -> AcpAsgiApp: - """Create an ASGI app serving an ACP agent over Streamable HTTP + WebSocket. - - Args: - agent_factory: Called once per connection with the bound - ``AgentSideConnection`` to produce a per-connection ``Agent``. - - Returns: - An :class:`AcpAsgiApp` ASGI 3.0 application. +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from functools import partial + +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response, StreamingResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket + +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 + +__all__ = ["create_asgi_app"] + + +def create_asgi_app(agent_factory: AgentFactory, *, path: str = ACP_ENDPOINT_PATH) -> Starlette: + """Create a Starlette app with one agent instance per connection. + + The app handles POST/GET/DELETE and WebSocket at ``path`` (default: /acp). + The path is relative to any parent mount. When mounting the app, enter its + lifespan from the parent lifespan too. """ - return AcpAsgiApp(AcpServer(agent_factory)) + server = AcpServer(agent_factory) + + @asynccontextmanager + async def lifespan(app: Starlette) -> AsyncIterator[None]: + try: + yield + finally: + await server.close() + + async def websocket(websocket: WebSocket) -> None: + await handle_websocket(agent_factory, websocket) + + return Starlette( + routes=[ + Route(path, partial(_post, server), methods=["POST"]), + Route(path, partial(_get, server), methods=["GET"]), + Route(path, partial(_delete, server), methods=["DELETE"]), + WebSocketRoute(path, websocket), + ], + lifespan=lifespan, + ) + + +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 await server.handle_post( + message, + content_type=request.headers.get("content-type"), + connection_id=request.headers.get(CONNECTION_ID_HEADER), + session_id=request.headers.get(SESSION_ID_HEADER), + ) + + +async def _get(server: AcpServer, request: Request) -> Response: + if request.headers.get("upgrade", "").lower() == "websocket": + return JSONResponse({"error": "WebSocket upgrade must use the ws scope"}, status_code=400) + accept = request.headers.get("accept", "") + if CONTENT_TYPE_SSE not in accept and "*/*" not in accept: + return JSONResponse({"error": "Accept must include text/event-stream"}, status_code=406) + connection_id = request.headers.get(CONNECTION_ID_HEADER) + if connection_id is None: + return JSONResponse({"error": "Missing connection id"}, status_code=400) + session_id = request.headers.get(SESSION_ID_HEADER) + error = server.validate_stream(connection_id=connection_id, session_id=session_id) + if error is not None: + return error + return StreamingResponse( + server.open_stream(connection_id=connection_id, session_id=session_id), + media_type=CONTENT_TYPE_SSE, + headers={"Cache-Control": "no-cache"}, + ) + + +async def _delete(server: AcpServer, request: Request) -> Response: + return await server.handle_delete(connection_id=request.headers.get(CONNECTION_ID_HEADER)) diff --git a/src/acp/http/protocol.py b/src/acp/http/protocol.py index c15fdb4..126173a 100644 --- a/src/acp/http/protocol.py +++ b/src/acp/http/protocol.py @@ -35,7 +35,7 @@ CONTENT_TYPE_JSON = "application/json" CONTENT_TYPE_SSE = "text/event-stream" -# Endpoint path used by docs/examples (the adapter itself is path-agnostic). +# Default endpoint path for the server and docs/examples. ACP_ENDPOINT_PATH = "/acp" INITIALIZE_METHOD = AGENT_METHODS["initialize"] diff --git a/src/acp/http/server.py b/src/acp/http/server.py index 6c4f00f..1f960e9 100644 --- a/src/acp/http/server.py +++ b/src/acp/http/server.py @@ -1,19 +1,19 @@ -"""Framework-agnostic Streamable HTTP server core (port of #155 server.ts + connection.ts). +"""Streamable HTTP connection management and ACP message routing. -:class:`AcpServer` owns an in-memory :class:`ConnectionRegistry`. For each +:class:`AcpServer` owns a dictionary of active HTTP connections. For each ``initialize`` POST it mints a connection, binds an ``AgentSideConnection`` to an -in-memory transport pair, and returns an ``Acp-Connection-Id``. Subsequent +HTTP transport, and returns an ``Acp-Connection-Id``. Subsequent server→client messages produced by the agent are fanned out to the correct SSE stream (connection-scoped or session-scoped) based on their ``sessionId`` / correlated request id. -The core exposes small, transport-neutral entry points: +HTTP handlers return Starlette responses directly: -* :meth:`AcpServer.handle_post` — returns a :class:`PostResult` (status + body). +* :meth:`AcpServer.handle_post` — returns a Starlette response. * :meth:`AcpServer.open_stream` — returns an async byte iterator of SSE frames. * :meth:`AcpServer.handle_delete` — terminates a connection. -The ASGI adapter in :mod:`acp.http.asgi` maps these onto ASGI messages. +The Starlette application in :mod:`acp.http.asgi` supplies requests and streams. """ from __future__ import annotations @@ -22,11 +22,14 @@ import contextlib import uuid from collections.abc import Callable -from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any +try: + from starlette.responses import JSONResponse, Response +except ImportError as exc: + raise ImportError("The HTTP server requires the 'http' extra: pip install agent-client-protocol[http]") from exc + from .._sse import serialize_sse_event, serialize_sse_keepalive -from .._transport import memory_transport_pair from ..agent.connection import AgentSideConnection from .protocol import ( CONNECTION_ID_HEADER, @@ -39,17 +42,14 @@ ) if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncGenerator from ..interfaces import Agent __all__ = [ "AcpServer", "AgentFactory", - "ConnectionRegistry", - "ConnectionState", "OutboundStream", - "PostResult", ] AgentFactory = Callable[[AgentSideConnection], "Agent"] @@ -65,15 +65,6 @@ INITIALIZE_TIMEOUT_SECONDS = 30.0 -@dataclass -class PostResult: - """Outcome of a POST request.""" - - status: int - body: dict[str, Any] | None = None - headers: dict[str, str] = field(default_factory=dict) - - class OutboundStream: """A backpressure-aware buffer for server→client messages. @@ -82,7 +73,7 @@ class OutboundStream: :meth:`push` *awaits* until the consumer drains rather than dropping the message — dropping a JSON-RPC response would permanently hang the peer's pending request. Awaiting propagates backpressure up to the agent's message - pump, mirroring the ``ReadableStream`` backpressure in the TypeScript SDK. + handlers, mirroring the ``ReadableStream`` backpressure in the TypeScript SDK. """ def __init__(self, *, capacity: int = 1024) -> None: @@ -123,7 +114,7 @@ def _try_put_sentinel(self) -> bool: return False return True - async def iterate(self) -> AsyncIterator[dict[str, Any]]: + async def iterate(self) -> AsyncGenerator[dict[str, Any], None]: while True: message = await self._queue.get() if message is None: @@ -131,198 +122,72 @@ async def iterate(self) -> AsyncIterator[dict[str, Any]]: yield message -class ConnectionState: - """Owns an ``AgentSideConnection`` bound to an in-memory transport pair. +class _HttpTransport: + """Receive POST messages and route agent output directly to HTTP/SSE. - The agent writes server→client messages onto the server end of the pair; a - pump task reads them and routes each to the connection-scoped stream or the - right session-scoped stream. + Only initialize returns in a POST body. New-session responses use the + connection stream; later responses follow their request's session route. + Requests and notifications from the agent carry their own sessionId. """ - def __init__(self, connection_id: str, agent_factory: AgentFactory, *, multiplex: bool = False) -> None: - self.connection_id = connection_id - # ``server_side`` is what the AgentSideConnection talks over; ``pump_side`` - # is what we read agent→client traffic from and inject client→agent on. - server_side, pump_side = memory_transport_pair() - self._pump_side = pump_side - self._agent_conn = AgentSideConnection(agent_factory, server_side, listening=True) + def __init__(self, initialize_id: Any) -> None: + self._initialize_id = message_id_key(initialize_id) + self.initialize_response: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future() + self._incoming: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() + self._closed = False self.connection_stream = OutboundStream() self.session_streams: dict[str, OutboundStream] = {} - # WebSocket mode: multiplex *all* agent→client traffic onto one stream - # (the single socket) instead of splitting across SSE streams. - self._multiplex: OutboundStream | None = OutboundStream() if multiplex else None - # Maps a request id -> sessionId, so responses to session-scoped client - # requests route back onto the right session stream. self._pending_routes: dict[str, str] = {} - # Request ids whose response should be captured (e.g. initialize) instead - # of being pushed to a stream. - self._response_waiters: dict[str, asyncio.Future[dict[str, Any]]] = {} - self._pump_task: asyncio.Task[None] | None = None - - def start(self) -> None: - self._pump_task = asyncio.ensure_future(self._pump()) - async def _pump(self) -> None: - try: - while True: - message = await self._pump_side.receive() - if message is None: - return - await self._route_outbound(message) - except asyncio.CancelledError: - return + async def receive(self) -> dict[str, Any] | None: + return await self._incoming.get() - async def _route_outbound(self, message: dict[str, Any]) -> None: - """Route an agent→client message to the correct SSE stream. - - Rules (matching the RFD): - - * A response to a session-*establishing* request (``session/new`` / - ``session/load`` — result carries a ``sessionId``) goes on the - **connection-scoped** stream, because the client does not yet have the - session-scoped stream open. We register the session so its stream can - be opened on the next GET. - * A response to an already-session-scoped client request routes onto that - session's stream (looked up via ``_pending_routes`` by request id). - * A server→client message carrying a ``sessionId`` in params (a - notification or request) routes onto that session's stream. - * Everything else goes on the connection-scoped stream. - """ - if is_response_message(message): - await self._route_response(message) - return - # Requests/notifications: route by sessionId in params if present. - if self._multiplex is not None: - await self._multiplex.push(message) - return + async def send(self, message: dict[str, Any]) -> None: + if self._closed: + raise ConnectionError("Transport closed") session_id = session_id_from_params(message.get("params")) - if session_id is not None and session_id in self.session_streams: - await self.session_streams[session_id].push(message) - return - await self.connection_stream.push(message) - - async def _route_response(self, message: dict[str, Any]) -> None: - key = message_id_key(message.get("id")) - # A captured response (e.g. initialize) resolves its waiter instead of - # being pushed to any stream. - if key is not None and key in self._response_waiters: - waiter = self._response_waiters.pop(key) - if not waiter.done(): - waiter.set_result(message) - return - # Register any newly-established session so unknown-session validation - # succeeds regardless of transport. - established = session_id_from_result(message.get("result")) - if established is not None: - self.ensure_session_stream(established) - routed = self._pending_routes.pop(key, None) if key is not None else None - if self._multiplex is not None: - await self._multiplex.push(message) - return - # session/new | session/load results (``established``) go on the - # connection-scoped stream; already-session-scoped responses route to the - # session stream recorded when the request came in. - if established is None and routed is not None and routed in self.session_streams: - await self.session_streams[routed].push(message) - return - await self.connection_stream.push(message) + if is_response_message(message): + key = message_id_key(message.get("id")) + if key == self._initialize_id and not self.initialize_response.done(): + self.initialize_response.set_result(message) + return + session_id = self._pending_routes.pop(key, None) if key is not None else None + established = session_id_from_result(message.get("result")) + if established is not None: + self.session_streams.setdefault(established, OutboundStream()) + # The client must learn the session ID before opening its stream. + session_id = None + stream = ( + self.session_streams.get(session_id, self.connection_stream) + if session_id is not None + else self.connection_stream + ) + await stream.push(message) async def deliver_to_agent(self, message: dict[str, Any]) -> None: - """Inject a client→server message into the agent connection.""" - # Track session-scoped client requests so their responses route back. + if self._closed: + raise ConnectionError("Transport closed") if "id" in message and "method" in message: session_id = session_id_from_params(message.get("params")) - if session_id is not None: - key = message_id_key(message["id"]) - if key is not None: - self._pending_routes[key] = session_id - await self._pump_side.send(message) - - async def request_response(self, message: dict[str, Any]) -> dict[str, Any]: - """Send a request to the agent and await its correlated response. - - Used for the ``initialize`` POST, which is the one request whose response - is returned synchronously in the HTTP body rather than over an SSE stream. - """ - key = message_id_key(message.get("id")) - loop = asyncio.get_running_loop() - future: asyncio.Future[dict[str, Any]] = loop.create_future() - if key is not None: - self._response_waiters[key] = future - await self.deliver_to_agent(message) - return await asyncio.wait_for(future, timeout=INITIALIZE_TIMEOUT_SECONDS) - - def ensure_session_stream(self, session_id: str) -> OutboundStream: - stream = self.session_streams.get(session_id) - if stream is None: - stream = OutboundStream() - self.session_streams[session_id] = stream - return stream - - def has_session(self, session_id: str) -> bool: - return session_id in self.session_streams - - async def iter_all_outbound(self) -> AsyncIterator[dict[str, Any]]: - """Iterate every agent→client message (WebSocket multiplex mode).""" - if self._multiplex is None: - msg = "iter_all_outbound requires a multiplex connection (WebSocket)" - raise RuntimeError(msg) - async for message in self._multiplex.iterate(): - yield message + key = message_id_key(message["id"]) + if session_id is not None and key is not None: + self._pending_routes[key] = session_id + self._incoming.put_nowait(dict(message)) async def close(self) -> None: - if self._pump_task is not None: - self._pump_task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await self._pump_task + if self._closed: + return + self._closed = True + self._incoming.put_nowait(None) + self.initialize_response.cancel() + self._pending_routes.clear() self.connection_stream.close() for stream in self.session_streams.values(): stream.close() - if self._multiplex is not None: - self._multiplex.close() - with contextlib.suppress(Exception): - await self._pump_side.close() - with contextlib.suppress(Exception): - await self._agent_conn.close() - - -class ConnectionRegistry: - """In-memory ``connectionId -> ConnectionState`` registry.""" - - def __init__(self) -> None: - self._connections: dict[str, ConnectionState] = {} - - def create(self, agent_factory: AgentFactory) -> ConnectionState: - connection_id = uuid.uuid4().hex - state = ConnectionState(connection_id, agent_factory) - state.start() - self._connections[connection_id] = state - return state - - def create_multiplex(self, agent_factory: AgentFactory) -> ConnectionState: - """Create a connection whose agent→client traffic is multiplexed onto one - stream (used by the WebSocket transport).""" - connection_id = uuid.uuid4().hex - state = ConnectionState(connection_id, agent_factory, multiplex=True) - state.start() - self._connections[connection_id] = state - return state - - def get(self, connection_id: str) -> ConnectionState | None: - return self._connections.get(connection_id) - - async def remove(self, connection_id: str) -> None: - state = self._connections.pop(connection_id, None) - if state is not None: - await state.close() - - async def close_all(self) -> None: - for connection_id in list(self._connections): - await self.remove(connection_id) class AcpServer: - """Framework-agnostic Streamable HTTP + WebSocket server core. + """Manage ACP HTTP connections and return Starlette responses. Args: agent_factory: Called once per connection with the bound @@ -330,16 +195,8 @@ class AcpServer: """ def __init__(self, agent_factory: AgentFactory) -> None: - self._agent_factory = agent_factory - self._registry = ConnectionRegistry() - - @property - def registry(self) -> ConnectionRegistry: - return self._registry - - def create_websocket_connection(self) -> ConnectionState: - """Create a new multiplexed connection for a WebSocket upgrade.""" - return self._registry.create_multiplex(self._agent_factory) + self.agent_factory = agent_factory + self._connections: dict[str, tuple[AgentSideConnection, _HttpTransport]] = {} # -- POST --------------------------------------------------------------- @@ -350,60 +207,69 @@ async def handle_post( content_type: str | None, connection_id: str | None, session_id: str | None, - ) -> PostResult: + ) -> Response: if content_type is None or not content_type.lower().startswith("application/json"): - return PostResult(415, {"error": "Content-Type must be application/json"}) + return JSONResponse({"error": "Content-Type must be application/json"}, status_code=415) if isinstance(message, list): - return PostResult(501, {"error": "Batch requests are not supported"}) + return JSONResponse({"error": "Batch requests are not supported"}, status_code=501) if not isinstance(message, dict): - return PostResult(400, {"error": "Invalid JSON-RPC message"}) + return JSONResponse({"error": "Invalid JSON-RPC message"}, status_code=400) if is_initialize_request(message): return await self._handle_initialize(message) if connection_id is None: - return PostResult(400, {"error": "Missing connection id"}) - state = self._registry.get(connection_id) - if state is None: - return PostResult(404, {"error": "Unknown connection id"}) + return JSONResponse({"error": "Missing connection id"}, status_code=400) + connection = self._connections.get(connection_id) + if connection is None: + return JSONResponse({"error": "Unknown connection id"}, status_code=404) + _, transport = connection method = message.get("method") if method_requires_session_header(method) and session_id is None: - return PostResult(400, {"error": "Missing session id header"}) - if session_id is not None and not state.has_session(session_id): + return JSONResponse({"error": "Missing session id header"}, status_code=400) + if session_id is not None and session_id not in transport.session_streams: # A session-scoped POST references an unknown session. - return PostResult(404, {"error": "Unknown session id"}) + return JSONResponse({"error": "Unknown session id"}, status_code=404) - await state.deliver_to_agent(message) - return PostResult(202) + await transport.deliver_to_agent(message) + return Response(status_code=202) - async def _handle_initialize(self, message: dict[str, Any]) -> PostResult: - state = self._registry.create(self._agent_factory) + async def _handle_initialize(self, message: dict[str, Any]) -> Response: + connection_id = uuid.uuid4().hex + transport = _HttpTransport(message.get("id")) + conn = AgentSideConnection(self.agent_factory, transport) + self._connections[connection_id] = (conn, transport) # Deliver initialize to the agent and await its response so we can return # the 200 body synchronously (initialize is the one blocking POST). If the # agent never responds (timeout) or errors, tear the just-created - # connection down instead of leaking its pump task + agent connection. + # connection down instead of leaking its agent connection. try: - response = await state.request_response(message) - except TimeoutError: - await self._registry.remove(state.connection_id) - return PostResult(504, {"error": "initialize timed out"}) + await transport.deliver_to_agent(message) + response = await asyncio.wait_for(transport.initialize_response, timeout=INITIALIZE_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + await self.handle_delete(connection_id=connection_id) + return JSONResponse({"error": "initialize timed out"}, status_code=504) except Exception: - await self._registry.remove(state.connection_id) - return PostResult(500, {"error": "initialize failed"}) - return PostResult(200, response, {CONNECTION_ID_HEADER: state.connection_id}) + await self.handle_delete(connection_id=connection_id) + return JSONResponse({"error": "initialize failed"}, status_code=500) + except asyncio.CancelledError: + await self.handle_delete(connection_id=connection_id) + raise + return JSONResponse(response, headers={CONNECTION_ID_HEADER: connection_id}) # -- GET / SSE ---------------------------------------------------------- - def validate_stream(self, *, connection_id: str | None, session_id: str | None) -> PostResult | None: - """Validate a GET SSE request. Returns an error PostResult, or None if OK.""" + def validate_stream(self, *, connection_id: str | None, session_id: str | None) -> Response | None: + """Validate a GET SSE request. Returns an error response, or None if OK.""" if connection_id is None: - return PostResult(400, {"error": "Missing connection id"}) - state = self._registry.get(connection_id) - if state is None: - return PostResult(404, {"error": "Unknown connection id"}) - if session_id is not None and not state.has_session(session_id): - return PostResult(404, {"error": "Unknown session id"}) + return JSONResponse({"error": "Missing connection id"}, status_code=400) + connection = self._connections.get(connection_id) + if connection is None: + return JSONResponse({"error": "Unknown connection id"}, status_code=404) + _, transport = connection + if session_id is not None and session_id not in transport.session_streams: + return JSONResponse({"error": "Unknown session id"}, status_code=404) return None async def open_stream( @@ -411,17 +277,18 @@ async def open_stream( *, connection_id: str, session_id: str | None, - ) -> AsyncIterator[bytes]: + ) -> AsyncGenerator[bytes, None]: """Yield SSE byte frames for a connection- or session-scoped stream. Emits a keepalive comment whenever the stream is idle for longer than :data:`SSE_KEEPALIVE_INTERVAL_SECONDS` so that idle-timeout intermediaries (proxies, load balancers) do not close an otherwise-healthy stream. """ - state = self._registry.get(connection_id) - if state is None: + connection = self._connections.get(connection_id) + if connection is None: return - stream = state.ensure_session_stream(session_id) if session_id is not None else state.connection_stream + _, transport = connection + stream = transport.session_streams[session_id] if session_id is not None else transport.connection_stream messages = stream.iterate() pending: asyncio.Task[dict[str, Any]] | None = None try: @@ -449,13 +316,16 @@ async def open_stream( # -- DELETE ------------------------------------------------------------- - async def handle_delete(self, *, connection_id: str | None) -> PostResult: + async def handle_delete(self, *, connection_id: str | None) -> Response: if connection_id is None: - return PostResult(400, {"error": "Missing connection id"}) - if self._registry.get(connection_id) is None: - return PostResult(404, {"error": "Unknown connection id"}) - await self._registry.remove(connection_id) - return PostResult(202) + return JSONResponse({"error": "Missing connection id"}, status_code=400) + connection = self._connections.pop(connection_id, None) + if connection is None: + return JSONResponse({"error": "Unknown connection id"}, status_code=404) + conn, _ = connection + await conn.close() + return Response(status_code=202) async def close(self) -> None: - await self._registry.close_all() + for connection_id in list(self._connections): + await self.handle_delete(connection_id=connection_id) diff --git a/src/acp/ws/server.py b/src/acp/ws/server.py index e5ff43f..f6179ea 100644 --- a/src/acp/ws/server.py +++ b/src/acp/ws/server.py @@ -1,79 +1,59 @@ -"""WebSocket server handling for the ASGI adapter (port of #155 ws-server.ts). - -On upgrade we create a fresh :class:`~acp.http.server.ConnectionState` (bound to -its own ``AgentSideConnection``), accept the socket with an ``Acp-Connection-Id`` -header, then pump JSON-RPC text frames both directions. All server→client -traffic (across the connection- and every session-scoped stream) is multiplexed -onto the single socket. On disconnect the connection and its sessions are torn -down. -""" +"""Bind a Starlette WebSocket to an ACP agent connection.""" from __future__ import annotations -import asyncio -import contextlib import json -from collections.abc import Callable +import uuid from typing import TYPE_CHECKING, Any +from starlette.websockets import WebSocket, WebSocketState + +from ..agent.connection import AgentSideConnection from ..http.protocol import CONNECTION_ID_HEADER if TYPE_CHECKING: - from ..http.server import AcpServer, ConnectionState - -__all__ = ["handle_asgi_websocket"] - - -async def handle_asgi_websocket( - server: AcpServer, - scope: dict[str, Any], - receive: Callable, - send: Callable, -) -> None: - """Handle an ASGI ``websocket`` scope by bridging it to a new ACP connection.""" - # Wait for the connect message. - message = await receive() - if message["type"] != "websocket.connect": - return - - state = server.create_websocket_connection() - await send({ - "type": "websocket.accept", - "headers": [(CONNECTION_ID_HEADER.lower().encode(), state.connection_id.encode())], - }) - - outbound_task = asyncio.ensure_future(_pump_outbound(state, send)) + from ..http.server import AgentFactory + +__all__ = ["handle_websocket"] + + +class _WebSocketTransport: + """Adapt Starlette's WebSocket to the message-level Transport interface.""" + + def __init__(self, websocket: WebSocket) -> None: + self._ws = websocket + + async def send(self, message: dict[str, 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: + while self._ws.client_state != WebSocketState.DISCONNECTED: + event = await self._ws.receive() + if event["type"] == "websocket.disconnect": + return None + if event.get("text") is None: + continue + try: + message = json.loads(event["text"]) + except json.JSONDecodeError: + continue + if isinstance(message, dict): + return message + return None + + async def close(self) -> None: + if self._ws.client_state == self._ws.application_state == WebSocketState.CONNECTED: + await self._ws.close() + + +async def handle_websocket(agent_factory: AgentFactory, websocket: WebSocket) -> None: + """Run one agent for the lifetime of the socket; disconnect cancels its work.""" + await websocket.accept(headers=[(CONNECTION_ID_HEADER.lower().encode(), uuid.uuid4().hex.encode())]) + transport = _WebSocketTransport(websocket) + conn = AgentSideConnection(agent_factory, transport, listening=False) try: - await _pump_inbound(state, receive) + await conn.listen() finally: - outbound_task.cancel() - with contextlib.suppress(asyncio.CancelledError, Exception): - await outbound_task - await server.registry.remove(state.connection_id) - - -async def _pump_inbound(state: ConnectionState, receive: Callable) -> None: - """Read client→server text frames and deliver them to the agent.""" - while True: - message = await receive() - msg_type = message["type"] - if msg_type == "websocket.disconnect": - return - if msg_type != "websocket.receive": - continue - text = message.get("text") - if text is None: - # Ignore binary frames. - continue - try: - payload = json.loads(text) - except json.JSONDecodeError: - continue - if isinstance(payload, dict): - await state.deliver_to_agent(payload) - - -async def _pump_outbound(state: ConnectionState, send: Callable) -> None: - """Forward all agent→client messages onto the socket as text frames.""" - async for message in state.iter_all_outbound(): - await send({"type": "websocket.send", "text": json.dumps(message, separators=(",", ":"))}) + await conn.close() diff --git a/tests/http/test_asgi.py b/tests/http/test_asgi.py new file mode 100644 index 0000000..b91940e --- /dev/null +++ b/tests/http/test_asgi.py @@ -0,0 +1,178 @@ +"""HTTP boundaries, mounting, and streaming cleanup of the Starlette app.""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import asynccontextmanager +from typing import Any + +import httpx +import pytest +from starlette.applications import Starlette +from starlette.routing import Mount +from websockets.asyncio.client import connect +from websockets.exceptions import InvalidStatus + +from acp.http.asgi import create_asgi_app +from acp.http.protocol import CONNECTION_ID_HEADER +from tests.conftest import TestAgent + +INITIALIZE = {"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {"protocolVersion": 1}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method", "headers", "body", "status"), + [ + ("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), + ("GET", {"Accept": "application/json"}, "", 406), + ("GET", {"Accept": "text/event-stream"}, "", 400), + ("GET", {"Accept": "text/event-stream", CONNECTION_ID_HEADER: "unknown"}, "", 404), + ("DELETE", {}, "", 400), + ("DELETE", {CONNECTION_ID_HEADER: "unknown"}, "", 404), + ], +) +async def test_http_errors(method: str, headers: dict[str, str], body: str, status: int) -> None: + app = create_asgi_app(lambda conn: TestAgent()) + async with ( + app.router.lifespan_context(app), + httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client, + ): + response = await client.request(method, "/acp", headers=headers, content=body) + assert response.status_code == status + assert "error" in response.json() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["PUT", "OPTIONS"]) +async def test_unsupported_methods_do_not_open_a_stream(method: str) -> None: + app = create_asgi_app(lambda conn: TestAgent()) + async with ( + app.router.lifespan_context(app), + httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client, + ): + response = await client.request(method, "/acp") + assert response.status_code == 405 + assert "allow" in response.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mounted", [False, True]) +@pytest.mark.parametrize("endpoint", [None, "/rpc"]) +async def test_direct_and_mounted_app_support_http_and_websocket( + mounted: bool, endpoint: str | None, serve_asgi +) -> None: + acp_app = ( + create_asgi_app(lambda conn: TestAgent()) + if endpoint is None + else create_asgi_app(lambda conn: TestAgent(), path=endpoint) + ) + assert isinstance(acp_app, Starlette) + + @asynccontextmanager + async def lifespan(app: Starlette): + async with acp_app.router.lifespan_context(acp_app): + yield + + app = Starlette(routes=[Mount("/agents", app=acp_app)], lifespan=lifespan) if mounted else acp_app + prefix = "/agents" if mounted else "" + path = prefix + (endpoint or "/acp") + server = await serve_asgi(app) + async with httpx.AsyncClient(base_url=f"http://{server.host}:{server.port}") as client: + response = await client.post(path, json=INITIALIZE) + assert response.status_code == 200 + assert response.json()["id"] == 0 + connection_id = response.headers[CONNECTION_ID_HEADER] + deleted = await client.delete(path, headers={CONNECTION_ID_HEADER: connection_id}) + assert deleted.status_code == 202 + assert deleted.content == b"" + assert (await client.delete(path, headers={CONNECTION_ID_HEADER: connection_id})).status_code == 404 + + async with connect(f"ws://{server.host}:{server.port}{path}") as websocket: + assert websocket.response is not None + assert CONNECTION_ID_HEADER in websocket.response.headers + await websocket.send(json.dumps(INITIALIZE)) + response_body = await asyncio.wait_for(websocket.recv(), timeout=1) + assert json.loads(response_body)["result"]["protocolVersion"] == 1 + + wrong_path = prefix + ("/acp" if endpoint else "/other") + for method in ("POST", "GET", "DELETE"): + assert (await client.request(method, wrong_path, json=INITIALIZE)).status_code == 404 + with pytest.raises(InvalidStatus) as exc: + async with connect(f"ws://{server.host}:{server.port}{wrong_path}"): + pytest.fail("WebSocket connected outside the configured path") + assert exc.value.response.status_code == 403 + + +@pytest.mark.asyncio +async def test_lifespan_closes_http_connections() -> None: + app = create_asgi_app(lambda conn: TestAgent()) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + async with app.router.lifespan_context(app): + response = await client.post("/acp", json=INITIALIZE) + connection_id = response.headers[CONNECTION_ID_HEADER] + response = await client.delete("/acp", headers={CONNECTION_ID_HEADER: connection_id}) + assert response.status_code == 404 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/acp", "/rpc"]) +async def test_sse_disconnect_releases_reader_before_reopening(path: str) -> None: + connections = [] + + def factory(conn): + connections.append(conn) + return TestAgent() + + app = create_asgi_app(factory, path=path) + async with app.router.lifespan_context(app): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response = await client.post(path, json=INITIALIZE) + connection_id = response.headers[CONNECTION_ID_HEADER] + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "method": "GET", + "path": path, + "root_path": "", + "query_string": b"", + "headers": [ + (b"accept", b"text/event-stream"), + (CONNECTION_ID_HEADER.lower().encode(), connection_id.encode()), + ], + } + disconnect = asyncio.Event() + started = asyncio.Event() + + async def receive() -> dict[str, Any]: + await disconnect.wait() + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] == "http.response.start": + assert message["status"] == 200 + started.set() + + first = asyncio.create_task(app(scope, receive, send)) + try: + await asyncio.wait_for(started.wait(), timeout=1) + disconnect.set() + await asyncio.wait_for(first, timeout=1) + finally: + first.cancel() + # A stale reader must not consume output intended for the reopened GET. + await connections[0].ext_notification("test", {"message": "after disconnect"}) + disconnect.clear() + frames = [] + + async def collect(message: dict[str, Any]) -> None: + if message["type"] == "http.response.body" and message.get("body"): + frames.append(message["body"]) + disconnect.set() + + await asyncio.wait_for(app(scope, receive, collect), timeout=1) + assert json.loads(frames[0].removeprefix(b"data: "))["params"] == {"message": "after disconnect"} diff --git a/tests/http/test_fixes.py b/tests/http/test_fixes.py index 2a8c307..1439934 100644 --- a/tests/http/test_fixes.py +++ b/tests/http/test_fixes.py @@ -68,6 +68,18 @@ async def test_outbound_stream_push_blocks_when_full() -> None: # -- Finding 3: SSE keepalive -------------------------------------------------- +@pytest.mark.asyncio +async def test_outbound_stream_close_unblocks_full_buffer() -> None: + stream = OutboundStream(capacity=1) + await stream.push({"n": 0}) + blocked = asyncio.create_task(stream.push({"n": 1})) + await asyncio.sleep(0) + assert not blocked.done() + stream.close() + await asyncio.wait_for(blocked, timeout=1) + assert [message async for message in stream.iterate()] == [] + + @pytest.mark.asyncio async def test_open_stream_emits_keepalive_when_idle(monkeypatch: pytest.MonkeyPatch) -> None: """An idle connection-scoped stream must emit periodic SSE keepalive frames.""" @@ -113,13 +125,50 @@ async def test_initialize_timeout_cleans_up_connection(monkeypatch: pytest.Monke connection_id=None, session_id=None, ) - assert result.status >= 500 + assert result.status_code >= 500 # No connection should remain registered after a failed initialize. - assert server.registry.get(result.headers.get(CONNECTION_ID_HEADER, "")) is None - assert _registry_size(server) == 0 + assert not server._connections await server.close() +@pytest.mark.asyncio +@pytest.mark.parametrize("shutdown", [False, True]) +async def test_interrupted_initialize_cleans_up_agent(shutdown: bool) -> None: + started = asyncio.Event() + stopped = asyncio.Event() + + class Agent(_SilentInitAgent): + async def initialize(self, protocol_version: int = 1, **kwargs: Any) -> Any: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + + server = AcpServer(lambda conn: Agent()) + request = asyncio.create_task( + server.handle_post( + {"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {"protocolVersion": 1}}, + content_type=CT_JSON, + connection_id=None, + session_id=None, + ) + ) + try: + await asyncio.wait_for(started.wait(), timeout=1) + if shutdown: + await asyncio.wait_for(server.close(), timeout=1) + else: + request.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(request, timeout=1) + assert stopped.is_set() + assert not server._connections + finally: + request.cancel() + await server.close() + + # -- Finding 4: HTTP client surfaces disconnect on stream EOF ------------------ @@ -159,10 +208,6 @@ def handler(request: httpx.Request) -> httpx.Response: # -- Helpers ------------------------------------------------------------------- -def _registry_size(server: AcpServer) -> int: - return len(server.registry._connections) # type: ignore[attr-defined] - - class _NoopAgent: def __init__(self) -> None: self._conn: Any = None diff --git a/tests/http/test_http_server.py b/tests/http/test_http_server.py index 92f9290..a0b6cf0 100644 --- a/tests/http/test_http_server.py +++ b/tests/http/test_http_server.py @@ -1,12 +1,14 @@ -"""Tests for the framework-agnostic AcpServer core (ported from server*.test.ts).""" +"""Tests for the AcpServer HTTP handlers (ported from server*.test.ts).""" from __future__ import annotations import asyncio +import json from typing import Any import pytest +from acp.exceptions import RequestError from acp.http.protocol import CONNECTION_ID_HEADER from acp.http.server import AcpServer from acp.schema import NewSessionResponse, PromptResponse @@ -62,7 +64,7 @@ async def test_post_wrong_content_type_returns_415() -> None: connection_id=None, session_id=None, ) - assert result.status == 415 + assert result.status_code == 415 await server.close() @@ -70,7 +72,7 @@ async def test_post_wrong_content_type_returns_415() -> None: 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) - assert result.status == 501 + assert result.status_code == 501 await server.close() @@ -83,10 +85,10 @@ async def test_initialize_creates_connection_and_returns_id() -> None: connection_id=None, session_id=None, ) - assert result.status == 200 + assert result.status_code == 200 assert CONNECTION_ID_HEADER in result.headers assert result.body is not None - assert result.body["id"] == 0 + assert json.loads(bytes(result.body))["id"] == 0 await server.close() @@ -99,7 +101,7 @@ async def test_missing_connection_id_returns_400() -> None: connection_id=None, session_id=None, ) - assert result.status == 400 + assert result.status_code == 400 await server.close() @@ -112,7 +114,7 @@ async def test_unknown_connection_id_returns_404() -> None: connection_id="nope", session_id=None, ) - assert result.status == 404 + assert result.status_code == 404 await server.close() @@ -139,7 +141,7 @@ async def test_session_new_result_on_connection_stream() -> None: connection_id=conn_id, session_id=None, ) - assert result.status == 202 + assert result.status_code == 202 await asyncio.sleep(0.1) joined = b"".join(frames).decode() assert '"sessionId":"sess-1"' in joined @@ -158,7 +160,7 @@ async def test_session_scoped_missing_session_header_returns_400() -> None: connection_id=conn_id, session_id=None, ) - assert result.status == 400 + assert result.status_code == 400 await server.close() @@ -185,7 +187,7 @@ async def test_prompt_streams_notification_on_session_stream() -> None: connection_id=conn_id, session_id="sess-1", ) - assert result.status == 202 + assert result.status_code == 202 await asyncio.sleep(0.15) session_joined = b"".join(session_frames).decode() # The agent_message_chunk notification is session-scoped. @@ -202,7 +204,7 @@ async def test_delete_terminates_connection() -> None: server = AcpServer(_agent_factory(_Agent())) conn_id = await _initialize(server) result = await server.handle_delete(connection_id=conn_id) - assert result.status == 202 + assert result.status_code == 202 # Subsequent use of the connection id 404s. follow = await server.handle_post( {"jsonrpc": "2.0", "id": 5, "method": "session/new", "params": {}}, @@ -210,7 +212,7 @@ async def test_delete_terminates_connection() -> None: connection_id=conn_id, session_id=None, ) - assert follow.status == 404 + assert follow.status_code == 404 await server.close() @@ -218,16 +220,80 @@ async def test_delete_terminates_connection() -> None: async def test_delete_missing_connection_id_returns_400() -> None: server = AcpServer(_agent_factory(_Agent())) result = await server.handle_delete(connection_id=None) - assert result.status == 400 + assert result.status_code == 400 await server.close() @pytest.mark.asyncio async def test_get_validation_errors() -> None: server = AcpServer(_agent_factory(_Agent())) - assert server.validate_stream(connection_id=None, session_id=None).status == 400 # type: ignore[union-attr] - assert server.validate_stream(connection_id="nope", session_id=None).status == 404 # type: ignore[union-attr] + assert server.validate_stream(connection_id=None, session_id=None).status_code == 400 # type: ignore[union-attr] + assert server.validate_stream(connection_id="nope", session_id=None).status_code == 404 # type: ignore[union-attr] conn_id = await _initialize(server) - assert server.validate_stream(connection_id=conn_id, session_id="ghost").status == 404 # type: ignore[union-attr] + assert server.validate_stream(connection_id=conn_id, session_id="ghost").status_code == 404 # type: ignore[union-attr] assert server.validate_stream(connection_id=conn_id, session_id=None) is None await server.close() + + +@pytest.mark.asyncio +async def test_concurrent_session_results_and_errors_stay_on_their_streams() -> None: + class Agent(_Agent): + async def new_session( + self, cwd: str | None = None, mcp_servers: Any = None, **kwargs: Any + ) -> NewSessionResponse: + return NewSessionResponse(session_id=str(cwd)) + + async def prompt(self, session_id: str, prompt: Any = None, **kwargs: Any) -> PromptResponse: + if session_id == "first": + raise RequestError(-32000, "prompt failed") + return PromptResponse(stop_reason="end_turn") + + server = AcpServer(lambda conn: Agent()) + connection_id = await _initialize(server) + connection_stream = server.open_stream(connection_id=connection_id, session_id=None) + try: + for request_id, session_id in enumerate(("first", "second")): + # Reuse initialize's id=0: its completed waiter must not intercept this. + await server.handle_post( + { + "jsonrpc": "2.0", + "id": request_id, + "method": "session/new", + "params": {"cwd": session_id, "mcpServers": []}, + }, + content_type=CT_JSON, + connection_id=connection_id, + session_id=None, + ) + frame = await asyncio.wait_for(anext(connection_stream), timeout=1) + assert json.loads(frame.removeprefix(b"data: "))["result"]["sessionId"] == session_id + + for request_id, session_id in enumerate(("first", "second"), start=2): + result = await server.handle_post( + { + "jsonrpc": "2.0", + "id": request_id, + "method": "session/prompt", + "params": {"sessionId": session_id, "prompt": []}, + }, + content_type=CT_JSON, + connection_id=connection_id, + session_id=session_id, + ) + assert result.status_code == 202 + + for request_id, session_id in enumerate(("first", "second"), start=2): + stream = server.open_stream(connection_id=connection_id, session_id=session_id) + try: + frame = await asyncio.wait_for(anext(stream), timeout=1) + response = json.loads(frame.removeprefix(b"data: ")) + assert response["id"] == request_id + if session_id == "first": + assert response["error"]["message"] == "prompt failed" + else: + assert response["result"]["stopReason"] == "end_turn" + finally: + await stream.aclose() + finally: + await connection_stream.aclose() + await server.close() diff --git a/tests/http/test_websocket.py b/tests/http/test_websocket.py index cf21cf5..bb88c07 100644 --- a/tests/http/test_websocket.py +++ b/tests/http/test_websocket.py @@ -7,13 +7,13 @@ from typing import Any import pytest +from starlette.websockets import WebSocket from websockets.asyncio.server import serve from acp.http.protocol import CONNECTION_ID_HEADER -from acp.http.server import AcpServer from acp.schema import NewSessionResponse, PromptResponse from acp.ws.client import create_websocket_stream -from acp.ws.server import handle_asgi_websocket +from acp.ws.server import handle_websocket from tests.conftest import TestAgent @@ -92,7 +92,7 @@ async def send_binary_then_text(ws: Any) -> None: class _FakeAsgiSocket: - """In-memory ASGI websocket double driving handle_asgi_websocket.""" + """In-memory ASGI websocket double driving handle_websocket.""" def __init__(self) -> None: self._incoming: asyncio.Queue[dict[str, Any]] = asyncio.Queue() @@ -133,24 +133,82 @@ async def _poll() -> dict[str, Any]: @pytest.mark.asyncio async def test_asgi_websocket_handshake_returns_connection_id() -> None: - server = AcpServer(lambda conn: _Agent()) socket = _FakeAsgiSocket() socket.client_connect() - handler = asyncio.ensure_future(handle_asgi_websocket(server, {"type": "websocket"}, socket.receive, socket.send)) + handler = asyncio.ensure_future( + handle_websocket(lambda conn: _Agent(), WebSocket({"type": "websocket"}, socket.receive, socket.send)) + ) await asyncio.sleep(0.05) header_names = [k for k, _ in socket.accepted_headers] assert CONNECTION_ID_HEADER.lower().encode() in header_names socket.client_disconnect() await asyncio.wait_for(handler, timeout=1) - await server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_asgi_websocket_cleans_up_running_prompt(disconnect: bool) -> None: + started = asyncio.Event() + stopped = asyncio.Event() + + class Agent(_Agent): + async def prompt(self, session_id: str, prompt: Any = None, **kwargs: Any) -> PromptResponse: + started.set() + try: + await asyncio.Event().wait() + finally: + stopped.set() + return PromptResponse(stop_reason="end_turn") + + socket = _FakeAsgiSocket() + socket.client_connect() + handler = asyncio.create_task( + handle_websocket(lambda conn: Agent(), WebSocket({"type": "websocket"}, socket.receive, socket.send)) + ) + try: + socket.client_send_text({ + "jsonrpc": "2.0", + "id": 1, + "method": "session/prompt", + "params": {"sessionId": "sess-ws", "prompt": []}, + }) + await asyncio.wait_for(started.wait(), timeout=1) + if disconnect: + socket.client_disconnect() + else: + handler.cancel() + await asyncio.wait_for(handler, timeout=1) + assert stopped.is_set() + finally: + socket.client_disconnect() + await asyncio.wait_for(handler, timeout=1) + + +@pytest.mark.asyncio +async def test_asgi_websocket_ignores_non_message_frames() -> None: + socket = _FakeAsgiSocket() + socket.client_connect() + for frame in ({"bytes": b"binary"}, {"text": "invalid json"}, {"text": "[]"}, {"text": "null"}): + socket._incoming.put_nowait({"type": "websocket.receive", **frame}) + handler = asyncio.create_task( + handle_websocket(lambda conn: _Agent(), WebSocket({"type": "websocket"}, socket.receive, socket.send)) + ) + try: + socket.client_send_text({"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {"protocolVersion": 1}}) + response = await socket.wait_for_send(lambda m: m.get("id") == 0) + assert "result" in response + finally: + socket.client_disconnect() + await asyncio.wait_for(handler, timeout=1) @pytest.mark.asyncio async def test_asgi_websocket_full_flow() -> None: - server = AcpServer(lambda conn: _Agent()) socket = _FakeAsgiSocket() socket.client_connect() - handler = asyncio.ensure_future(handle_asgi_websocket(server, {"type": "websocket"}, socket.receive, socket.send)) + handler = asyncio.ensure_future( + handle_websocket(lambda conn: _Agent(), WebSocket({"type": "websocket"}, socket.receive, socket.send)) + ) await asyncio.sleep(0.05) socket.client_send_text({"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {"protocolVersion": 1}}) @@ -179,4 +237,3 @@ async def test_asgi_websocket_full_flow() -> None: socket.client_disconnect() await asyncio.wait_for(handler, timeout=1) - await server.close() diff --git a/uv.lock b/uv.lock index 30e9799..18190d3 100644 --- a/uv.lock +++ b/uv.lock @@ -18,6 +18,7 @@ dependencies = [ [package.optional-dependencies] http = [ { name = "httpx", extra = ["http2"] }, + { name = "starlette" }, { name = "websockets" }, ] logfire = [ @@ -38,6 +39,7 @@ dev = [ { name = "pytest-asyncio" }, { name = "python-dotenv" }, { name = "ruff" }, + { name = "starlette" }, { name = "tox-uv" }, { name = "ty" }, { name = "uvicorn" }, @@ -51,6 +53,7 @@ requires-dist = [ { name = "opentelemetry-sdk", marker = "extra == 'logfire'", specifier = ">=1.28.0" }, { name = "pydantic", specifier = ">=2.7" }, { name = "pydantic-core", specifier = ">=2.18.1" }, + { name = "starlette", marker = "extra == 'http'", specifier = ">=0.49.3" }, { name = "websockets", marker = "extra == 'http'", specifier = ">=12.0" }, ] provides-extras = ["logfire", "http"] @@ -68,6 +71,7 @@ dev = [ { name = "pytest-asyncio", specifier = ">=0.21.0" }, { name = "python-dotenv", specifier = ">=1.1.1" }, { name = "ruff", specifier = ">=0.11.5" }, + { name = "starlette", specifier = ">=0.49.3" }, { name = "tox-uv", specifier = ">=1.11.3" }, { name = "ty", specifier = ">=0.0.1a16" }, { name = "uvicorn", specifier = ">=0.30" }, @@ -1461,6 +1465,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "starlette" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b5/b4/205b0d5241d934e8add0c38aa924c4f9fb7330834ff11e5444db964ec3f9/starlette-1.6.0.tar.gz", hash = "sha256:d4e3ac5e546444960c710297a3c9fc3f7ebae1b7e963f3d36173b49da535be9b", size = 2716969, upload-time = "2026-08-08T18:27:57.512Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/cb/6a6a47d5b464bd08695d254f3da6e7986cc70c9fa5d778eda57538edfe56/starlette-1.6.0-py3-none-any.whl", hash = "sha256:a86dd39d14bb45f85a3d18525215a9ef0cfd1f192ac793220e72598c90335f0c", size = 75969, upload-time = "2026-08-08T18:27:56.196Z" }, +] + [[package]] name = "tomli" version = "2.3.0" From 4fddab307ac602946270233ee67d35c6af74ca57 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Fri, 11 Sep 2026 11:30:44 +0800 Subject: [PATCH 2/2] fix: support loading sessions over web transports --- docs/web-transport.md | 28 ++++++++++- src/acp/http/client.py | 31 ++++++++++-- src/acp/http/protocol.py | 2 + src/acp/http/server.py | 43 ++++++++++++++--- tests/http/test_http_client.py | 32 +++++++++++++ tests/http/test_http_server.py | 43 ++++++++++++++++- tests/http/test_loopback.py | 88 +++++++++++++++++++++++++++++++++- 7 files changed, 251 insertions(+), 16 deletions(-) diff --git a/docs/web-transport.md b/docs/web-transport.md index a6fb166..20c000b 100644 --- a/docs/web-transport.md +++ b/docs/web-transport.md @@ -55,6 +55,31 @@ header, then opens the connection-scoped SSE stream. When a new `sessionId` appears it opens that session-scoped stream too. A single SSE attempt is made per stream; reconnect/retry is the caller's responsibility (v1 of the RFD). +### Loading an existing session + +Both HTTP and WebSocket support `load_session()` when the agent advertises +`loadSession` and implements session persistence: + +```python +init = await conn.initialize(protocol_version=1) +if init.agent_capabilities.load_session: + await conn.load_session(session_id="saved-session-id", cwd="/workspace", mcp_servers=[]) + await conn.prompt(session_id="saved-session-id", prompt=[...]) +``` + +For HTTP, history replay and the load response use the connection SSE stream. +The client correlates the response with the session ID from the load request, +then opens the session SSE stream for further prompts and agent callbacks. This +also works with an empty history: a load response need not contain `sessionId`. +Replay is consumed as it arrives, so histories larger than the SSE buffer do not +wait for a session stream to open. WebSocket uses its existing bidirectional +connection for both replay and subsequent messages. + +A failed load returns its JSON-RPC error on the connection stream and can be +retried. The server removes streams provisioned only for failed loads, while +preserving established sessions and overlapping loads. It does not change the +agent's load response or automatically enable the agent's `loadSession` capability. + ## Server The server uses Starlette for HTTP requests, responses, routing, streaming, @@ -122,12 +147,13 @@ has one incoming queue and one SSE buffer per stream. The incoming queue lets POST return `202` while the agent handles the request. Output goes directly to the relevant SSE buffer; there is no intermediate transport pair or pump task. -HTTP output needs three routing rules: +HTTP output follows these routing rules: | Message | Destination | Why | | --- | --- | --- | | `initialize` response | POST body, via one Future | Establishes the connection before GET streams open | | Response containing a new `sessionId` | Connection SSE stream | The client needs the ID before it can open the session stream | +| `session/load` replay and response | Connection SSE stream | Replay precedes the response; the client gets the session ID from the original request | | Other messages | Session SSE stream when known, otherwise connection stream | Responses use their request's recorded session; requests/notifications carry `sessionId` | `OutboundStream` retains a bounded buffer, backpressure, and close handling. diff --git a/src/acp/http/client.py b/src/acp/http/client.py index b224167..77381f8 100644 --- a/src/acp/http/client.py +++ b/src/acp/http/client.py @@ -28,8 +28,11 @@ CONNECTION_ID_HEADER, CONTENT_TYPE_JSON, CONTENT_TYPE_SSE, + LOAD_SESSION_METHOD, SESSION_ID_HEADER, is_initialize_request, + is_response_message, + message_id_key, method_requires_session_header, session_id_from_message, ) @@ -78,6 +81,7 @@ def __init__( self._inbox: asyncio.Queue[Any] = asyncio.Queue() self._stream_tasks: set[asyncio.Task[None]] = set() self._session_streams: set[str] = set() + self._pending_loads: dict[str, str] = {} # -- Transport protocol ------------------------------------------------- @@ -87,7 +91,16 @@ async def send(self, message: dict[str, Any]) -> None: if is_initialize_request(message): await self._send_initialize(message) return - await self._send_post(message) + key = message_id_key(message.get("id")) + session_id = session_id_from_message(message) + if message.get("method") == LOAD_SESSION_METHOD and key is not None and session_id is not None: + self._pending_loads[key] = session_id + try: + await self._send_post(message) + except BaseException: + if key is not None: + self._pending_loads.pop(key, None) + raise async def receive(self) -> dict[str, Any] | None: item = await self._inbox.get() @@ -99,6 +112,7 @@ async def close(self) -> None: if self._closed: return self._closed = True + self._pending_loads.clear() for task in list(self._stream_tasks): task.cancel() for task in list(self._stream_tasks): @@ -148,7 +162,7 @@ async def _send_post(self, message: dict[str, Any]) -> None: # Some servers may answer initialize-like 200 bodies; for 200 with a body enqueue it. if response.status_code == 200 and response.content: with contextlib.suppress(Exception): - self._inbox.put_nowait(response.json()) + self._handle_incoming(response.json()) def _open_stream(self, *, session_id: str | None) -> None: if self._closed: @@ -197,13 +211,20 @@ def _on_stream_closed(self, session_id: str | None) -> None: self._session_streams.discard(session_id) return if not self._closed: + self._pending_loads.clear() self._inbox.put_nowait(_EOF) def _handle_incoming(self, message: dict[str, Any]) -> None: - # Open a session-scoped stream when any message carries a new sessionId - # (e.g. a session/new or session/load result on the connection stream). + # Load responses may be empty or null; the session ID is in the request. + if is_response_message(message): + key = message_id_key(message.get("id")) + loaded = self._pending_loads.pop(key, None) if key is not None else None + if loaded is not None and "result" in message: + self._open_stream(session_id=loaded) session_id = session_id_from_message(message) - if session_id is not None and session_id not in self._session_streams: + # Replay stays on the connection stream until load succeeds. Avoid + # opening a session GET that a failed load would immediately tear down. + if session_id is not None and session_id not in self._pending_loads.values(): self._open_stream(session_id=session_id) self._inbox.put_nowait(message) diff --git a/src/acp/http/protocol.py b/src/acp/http/protocol.py index 126173a..e61476a 100644 --- a/src/acp/http/protocol.py +++ b/src/acp/http/protocol.py @@ -17,6 +17,7 @@ "CONTENT_TYPE_JSON", "CONTENT_TYPE_SSE", "INITIALIZE_METHOD", + "LOAD_SESSION_METHOD", "SESSION_ID_HEADER", "is_initialize_request", "is_response_message", @@ -39,6 +40,7 @@ ACP_ENDPOINT_PATH = "/acp" INITIALIZE_METHOD = AGENT_METHODS["initialize"] +LOAD_SESSION_METHOD = AGENT_METHODS["session_load"] # Agent methods that operate on an *already-established* session and therefore # require the ``Acp-Session-Id`` header on POST + session-scoped routing of their diff --git a/src/acp/http/server.py b/src/acp/http/server.py index 1f960e9..1138ea4 100644 --- a/src/acp/http/server.py +++ b/src/acp/http/server.py @@ -33,6 +33,7 @@ from ..agent.connection import AgentSideConnection from .protocol import ( CONNECTION_ID_HEADER, + LOAD_SESSION_METHOD, is_initialize_request, is_response_message, message_id_key, @@ -138,6 +139,8 @@ def __init__(self, initialize_id: Any) -> None: self.connection_stream = OutboundStream() self.session_streams: dict[str, OutboundStream] = {} self._pending_routes: dict[str, str] = {} + self._pending_loads: dict[str, str] = {} + self._provisional_sessions: set[str] = set() async def receive(self) -> dict[str, Any] | None: return await self._incoming.get() @@ -151,12 +154,11 @@ async def send(self, message: dict[str, Any]) -> None: if key == self._initialize_id and not self.initialize_response.done(): self.initialize_response.set_result(message) return - session_id = self._pending_routes.pop(key, None) if key is not None else None - established = session_id_from_result(message.get("result")) - if established is not None: - self.session_streams.setdefault(established, OutboundStream()) - # The client must learn the session ID before opening its stream. - session_id = None + session_id = self._route_response(message, key) + # Replay and load responses share the connection stream, including on + # reload. This preserves replay order and never waits for a session GET. + if session_id in self._pending_loads.values(): + session_id = None stream = ( self.session_streams.get(session_id, self.connection_stream) if session_id is not None @@ -164,6 +166,24 @@ async def send(self, message: dict[str, Any]) -> None: ) await stream.push(message) + def _route_response(self, message: dict[str, Any], key: str | None) -> str | None: + loaded = self._pending_loads.pop(key, None) if key is not None else None + if loaded is not None: + if "result" in message: + self._provisional_sessions.discard(loaded) + elif loaded in self._provisional_sessions and loaded not in self._pending_loads.values(): + self._provisional_sessions.remove(loaded) + self.session_streams.pop(loaded).close() + return None + session_id = self._pending_routes.pop(key, None) if key is not None else None + established = session_id_from_result(message.get("result")) + if established is not None: + self.session_streams.setdefault(established, OutboundStream()) + self._provisional_sessions.discard(established) + # The client must learn the session ID before opening its stream. + return None + return session_id + async def deliver_to_agent(self, message: dict[str, Any]) -> None: if self._closed: raise ConnectionError("Transport closed") @@ -171,7 +191,14 @@ async def deliver_to_agent(self, message: dict[str, Any]) -> None: session_id = session_id_from_params(message.get("params")) key = message_id_key(message["id"]) if session_id is not None and key is not None: - self._pending_routes[key] = session_id + if message["method"] == LOAD_SESSION_METHOD: + self._pending_loads[key] = session_id + if session_id not in self.session_streams: + # Allow clients to open a GET as soon as replay starts. + self.session_streams[session_id] = OutboundStream() + self._provisional_sessions.add(session_id) + else: + self._pending_routes[key] = session_id self._incoming.put_nowait(dict(message)) async def close(self) -> None: @@ -181,6 +208,8 @@ async def close(self) -> None: self._incoming.put_nowait(None) self.initialize_response.cancel() self._pending_routes.clear() + self._pending_loads.clear() + self._provisional_sessions.clear() self.connection_stream.close() for stream in self.session_streams.values(): stream.close() diff --git a/tests/http/test_http_client.py b/tests/http/test_http_client.py index cacb296..18578f6 100644 --- a/tests/http/test_http_client.py +++ b/tests/http/test_http_client.py @@ -211,3 +211,35 @@ def handler(request: httpx.Request) -> httpx.Response: finally: await transport.close() await client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("result", [{}, None]) +async def test_load_opens_session_stream_from_request_id(result: Any) -> None: + server = FakeServer() + transport, client = _make_transport(server) + try: + await transport.send({"jsonrpc": "2.0", "id": 0, "method": "initialize", "params": {}}) + await asyncio.wait_for(transport.receive(), timeout=1) + await transport.send({ + "jsonrpc": "2.0", + "id": "load-1", + "method": "session/load", + "params": {"sessionId": "saved", "cwd": "/", "mcpServers": []}, + }) + replay = {"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "saved"}} + server.push_conn(replay) + assert await asyncio.wait_for(transport.receive(), timeout=1) == replay + await asyncio.sleep(0) + assert "saved" not in server.session_streams + + # A standard load response carries no sessionId, including null results. + response = {"jsonrpc": "2.0", "id": "load-1", "result": result} + server.push_conn(response) + assert await asyncio.wait_for(transport.receive(), timeout=1) == response + live = {"jsonrpc": "2.0", "id": 2, "result": {"stopReason": "end_turn"}} + server.push_session("saved", live) + assert await asyncio.wait_for(transport.receive(), timeout=1) == live + finally: + await transport.close() + await client.aclose() diff --git a/tests/http/test_http_server.py b/tests/http/test_http_server.py index a0b6cf0..382579e 100644 --- a/tests/http/test_http_server.py +++ b/tests/http/test_http_server.py @@ -10,7 +10,7 @@ from acp.exceptions import RequestError from acp.http.protocol import CONNECTION_ID_HEADER -from acp.http.server import AcpServer +from acp.http.server import AcpServer, _HttpTransport from acp.schema import NewSessionResponse, PromptResponse from tests.conftest import TestAgent @@ -297,3 +297,44 @@ async def prompt(self, session_id: str, prompt: Any = None, **kwargs: Any) -> Pr finally: await connection_stream.aclose() await server.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("first_succeeds", [False, True]) +@pytest.mark.parametrize("second_succeeds", [False, True]) +async def test_overlapping_loads_preserve_successful_session_streams( + first_succeeds: bool, second_succeeds: bool +) -> None: + transport = _HttpTransport(0) + connection_stream = transport.connection_stream.iterate() + try: + for request_id in (1, 2): + await transport.deliver_to_agent({ + "jsonrpc": "2.0", + "id": request_id, + "method": "session/load", + "params": {"sessionId": "saved", "cwd": "/", "mcpServers": []}, + }) + # A client may attach its session GET before load completes. + session_stream = transport.session_streams["saved"] + for request_id, succeeds in enumerate((first_succeeds, second_succeeds), start=1): + response = {"jsonrpc": "2.0", "id": request_id} + response.update({"result": {}} if succeeds else {"error": {"code": -32000, "message": "load failed"}}) + await transport.send(response) + assert await asyncio.wait_for(anext(connection_stream), timeout=1) == response + if request_id == 1: + assert transport.session_streams["saved"] is session_stream + + if first_succeeds or second_succeeds: + assert transport.session_streams["saved"] is session_stream + live = {"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "saved"}} + await transport.send(live) + messages = session_stream.iterate() + assert await asyncio.wait_for(anext(messages), timeout=1) == live + await messages.aclose() + else: + assert "saved" not in transport.session_streams + assert [message async for message in session_stream.iterate()] == [] + finally: + await connection_stream.aclose() + await transport.close() diff --git a/tests/http/test_loopback.py b/tests/http/test_loopback.py index 07b863d..7d8b8fe 100644 --- a/tests/http/test_loopback.py +++ b/tests/http/test_loopback.py @@ -12,10 +12,19 @@ import pytest -from acp import connect_to_agent +from acp import RequestError, connect_to_agent from acp.http.asgi import create_asgi_app from acp.http.client import create_http_stream -from acp.schema import InitializeResponse, NewSessionResponse, PromptResponse, RequestPermissionResponse +from acp.schema import ( + AgentCapabilities, + AgentMessageChunk, + InitializeResponse, + LoadSessionResponse, + NewSessionResponse, + PromptResponse, + RequestPermissionResponse, + TextContentBlock, +) from acp.ws.client import create_websocket_stream from tests.conftest import TestAgent, TestClient @@ -129,3 +138,78 @@ async def test_ws_loopback_prompt_streams_and_permission(serve_asgi) -> None: finally: await conn.close() await transport.close() + + +class _LoadingAgent(_LoopbackAgent): + def __init__(self, history_size: int) -> None: + super().__init__() + self.history_size = history_size + self.fail_load = False + self.ask_permission = True + + async def initialize(self, protocol_version: int = 1, **kwargs: Any) -> InitializeResponse: + return InitializeResponse(protocol_version=1, agent_capabilities=AgentCapabilities(load_session=True)) + + async def load_session(self, cwd: str, session_id: str, **kwargs: Any) -> LoadSessionResponse: + for i in range(self.history_size): + await self._conn.session_update( + session_id=session_id, + update=AgentMessageChunk(content=TextContentBlock(text=f"history-{i}")), + ) + if self.fail_load: + raise RequestError(-32000, "load failed") + return LoadSessionResponse() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protocol", ["http", "ws"]) +@pytest.mark.parametrize("history_size", [0, 1100]) +async def test_load_session_replays_history_and_supports_prompt(protocol: str, history_size: int, serve_asgi) -> None: + agent = _LoadingAgent(history_size) + server = await serve_asgi(_make_app(agent)) + transport = ( + create_http_stream(server.http_url) if protocol == "http" else await create_websocket_stream(server.ws_url) + ) + client = _CapturingClient() + conn = connect_to_agent(client, transport) + try: + init = await conn.initialize(protocol_version=1) + assert init.agent_capabilities.load_session + # Load an existing ID without first calling session/new, then reload it. + for _ in range(2): + client.updates.clear() + loaded = await asyncio.wait_for(conn.load_session(cwd="/", session_id="saved-session"), timeout=10) + assert loaded == LoadSessionResponse() + assert [update.content.text for update in client.updates] == [f"history-{i}" for i in range(history_size)] + result = await asyncio.wait_for(conn.prompt(session_id="saved-session", prompt=[]), timeout=10) + assert result.stop_reason == "end_turn" + assert client.permission_requested + assert client.updates[-1].content.text == "hello" + finally: + await conn.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("protocol", ["http", "ws"]) +async def test_failed_load_can_retry_and_preserves_existing_session(protocol: str, serve_asgi) -> None: + agent = _LoadingAgent(1) + server = await serve_asgi(_make_app(agent)) + transport = ( + create_http_stream(server.http_url) if protocol == "http" else await create_websocket_stream(server.ws_url) + ) + conn = connect_to_agent(_CapturingClient(), transport) + try: + await conn.initialize(protocol_version=1) + agent.fail_load = True + with pytest.raises(RequestError, match="load failed"): + await asyncio.wait_for(conn.load_session(cwd="/", session_id="saved-session"), timeout=5) + agent.fail_load = False + agent.history_size = 0 + await asyncio.wait_for(conn.load_session(cwd="/", session_id="saved-session"), timeout=5) + agent.fail_load = True + with pytest.raises(RequestError, match="load failed"): + await asyncio.wait_for(conn.load_session(cwd="/", session_id="saved-session"), timeout=5) + result = await asyncio.wait_for(conn.prompt(session_id="saved-session", prompt=[]), timeout=5) + assert result.stop_reason == "end_turn" + finally: + await conn.close()