fix: route session/load resume streams correctly

Author: federicociner

src/server-session-sse.test.ts

@@ -5,11 +5,18 @@ import {
   HEADER_SESSION_ID,
   JSON_MIME_TYPE,
 } from "./protocol.js";
+import { PROTOCOL_VERSION } from "./schema/index.js";
 import { parseSseStream } from "./sse.js";
 import { TestAgent } from "./test-support/test-agent.js";
 import { startTestServer } from "./test-support/test-http-server.js";
 
-import type { AgentSideConnection } from "./acp.js";
+import type {
+  AgentSideConnection,
+  InitializeRequest,
+  InitializeResponse,
+  LoadSessionRequest,
+  LoadSessionResponse,
+} from "./acp.js";
 import type { AnyMessage } from "./jsonrpc.js";
 
 const initializeRequest = {
@@ -57,6 +64,49 @@ function createForkRequest(id: number, sessionId: string) {
   };
 }
 
+function createLoadSessionRequest(id: number, sessionId: string) {
+  return {
+    jsonrpc: "2.0",
+    id,
+    method: "session/load",
+    params: {
+      cwd: "/tmp",
+      mcpServers: [],
+      sessionId,
+    },
+  };
+}
+
+class LoadSessionAgent extends TestAgent {
+  constructor(private readonly agentConnection: AgentSideConnection) {
+    super(agentConnection);
+  }
+
+  initialize(_params: InitializeRequest): Promise<InitializeResponse> {
+    return Promise.resolve({
+      protocolVersion: PROTOCOL_VERSION,
+      agentCapabilities: {
+        loadSession: true,
+      },
+    });
+  }
+
+  async loadSession(params: LoadSessionRequest): Promise<LoadSessionResponse> {
+    await this.agentConnection.sessionUpdate({
+      sessionId: params.sessionId,
+      update: {
+        sessionUpdate: "agent_message_chunk",
+        content: {
+          type: "text",
+          text: "replayed-session-history",
+        },
+      },
+    });
+
+    return {};
+  }
+}
+
 describe("AcpServer session SSE", () => {
   it("streams prompt updates and responses on the session SSE stream", async () => {
     const server = await startTestServer(
@@ -203,6 +253,58 @@ describe("AcpServer session SSE", () => {
     }
   });
 
+  it("routes session/load replay updates to session SSE and final response to connection SSE", async () => {
+    const server = await startTestServer(
+      (conn: AgentSideConnection) => new LoadSessionAgent(conn),
+    );
+
+    try {
+      const connectionId = await initialize(server.url);
+      const sessionId = "existing-session";
+      const connectionSse = await openConnectionSse(server.url, connectionId);
+      const sessionSse = await openSessionSse(
+        server.url,
+        connectionId,
+        sessionId,
+      );
+      const accepted = await postJson(
+        server.url,
+        createLoadSessionRequest(3, sessionId),
+        {
+          [HEADER_CONNECTION_ID]: connectionId,
+          [HEADER_SESSION_ID]: sessionId,
+        },
+      );
+
+      expect(sessionSse.status).toBe(200);
+      expect(accepted.status).toBe(202);
+      expect(await readSseMessages(sessionSse, 1)).toMatchObject([
+        {
+          jsonrpc: "2.0",
+          method: "session/update",
+          params: {
+            sessionId,
+            update: {
+              sessionUpdate: "agent_message_chunk",
+              content: {
+                text: "replayed-session-history",
+              },
+            },
+          },
+        },
+      ]);
+      expect(await readSseMessages(connectionSse, 1)).toMatchObject([
+        {
+          jsonrpc: "2.0",
+          id: 3,
+          result: {},
+        },
+      ]);
+    } finally {
+      await server.close();
+    }
+  });
+
   it("replays buffered session messages when session SSE attaches after prompt", async () => {
     const server = await startTestServer();
 

src/server.test.ts

@@ -176,7 +176,7 @@ describe("AcpServer", () => {
     }
   });
 
-  it("rejects session-scoped GETs for unknown sessions", async () => {
+  it("opens session-scoped GETs for sessions without local streams", async () => {
     const server = await startTestServer();
 
     try {
@@ -187,7 +187,7 @@ describe("AcpServer", () => {
         globalThis.crypto.randomUUID(),
       );
 
-      expect(response.status).toBe(404);
+      expect(response.status).toBe(200);
     } finally {
       await server.close();
     }

src/server.ts

@@ -14,6 +14,7 @@ import {
   isRequestMessage,
   isResponseMessage,
 } from "./jsonrpc.js";
+import { AGENT_METHODS } from "./schema/index.js";
 import { serializeSseEvent, serializeSseKeepAlive } from "./sse.js";
 import { handleWebSocketConnection } from "./ws-server.js";
 import type { WebSocketServerSocket } from "./ws-server.js";
@@ -179,12 +180,7 @@ export class AcpServer {
 
     const sessionId = req.headers.get(HEADER_SESSION_ID);
     if (sessionId) {
-      const sessionStream = connection.sessionStreams.get(sessionId);
-      if (!sessionStream) {
-        return textResponse("Unknown Acp-Session-Id", 404);
-      }
-
-      return sseResponse(sessionStream.subscribe());
+      return sseResponse(connection.ensureSession(sessionId).subscribe());
     }
 
     return sseResponse(connection.connectionStream.subscribe());
@@ -336,7 +332,10 @@ async function forwardClientRequest(
   const key = messageIdKey(message.id);
 
   if (key) {
-    connection.pendingRoutes.set(key, route.value);
+    connection.pendingRoutes.set(
+      key,
+      pendingResponseRoute(message, route.value),
+    );
   }
 
   await writeInbound(connection, message);
@@ -359,6 +358,13 @@ async function forwardClientNotification(
   return { ok: true };
 }
 
+function pendingResponseRoute(
+  message: ClientRequestMessage,
+  route: ResponseRoute,
+): ResponseRoute {
+  return message.method === AGENT_METHODS.session_load ? "connection" : route;
+}
+
 function determineRoute(
   message: ClientRequestMessage,
   headers: Headers,