fix: support loading sessions over web transports

Author: frostming

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.

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)
 

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

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,27 +154,51 @@ 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
             else self.connection_stream
         )
         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")
         if "id" in message and "method" in message:
             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()

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

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