fix: Align HTTP session routing with RFD

Author: federicociner

src/connection.ts

@@ -91,6 +91,7 @@ export class ConnectionState {
   readonly allOutbound = new OutboundStream();
   readonly sessionStreams = new Map<string, OutboundStream>();
   readonly pendingRoutes = new Map<string, ResponseRoute>();
+  readonly clientResponseRoutes = new Map<string, ResponseRoute>();
 
   private hasStartedRouter = false;
   private outboundReader: ReadableStreamDefaultReader<AnyMessage> | undefined;
@@ -163,6 +164,7 @@ export class ConnectionState {
 
     this.sessionStreams.clear();
     this.pendingRoutes.clear();
+    this.clientResponseRoutes.clear();
 
     await Promise.allSettled([
       this.inboundTx.close(),
@@ -231,13 +233,29 @@ export class ConnectionState {
   private routeOutboundRequestOrNotification(message: AnyMessage): void {
     const sessionId = sessionIdFromMessageParams(message);
     if (sessionId) {
+      this.trackClientResponseRoute(message, { session: sessionId });
       this.ensureSession(sessionId).push(message);
       return;
     }
 
+    this.trackClientResponseRoute(message, "connection");
     this.connectionStream.push(message);
   }
 
+  private trackClientResponseRoute(
+    message: AnyMessage,
+    route: ResponseRoute,
+  ): void {
+    if (!("id" in message) || !("method" in message)) {
+      return;
+    }
+
+    const key = messageIdKey(message.id);
+    if (key) {
+      this.clientResponseRoutes.set(key, route);
+    }
+  }
+
   private pushToRoute(route: ResponseRoute, message: AnyMessage): void {
     if (route === "connection") {
       this.connectionStream.push(message);

src/http-stream.test.ts

@@ -59,6 +59,48 @@ const promptRequest = {
   },
 } satisfies AnyMessage;
 
+const loadSessionRequest = {
+  jsonrpc: "2.0",
+  id: 3,
+  method: "session/load",
+  params: {
+    cwd: "/tmp",
+    mcpServers: [],
+    sessionId: "existing-session",
+  },
+} satisfies AnyMessage;
+
+const permissionRequest = {
+  jsonrpc: "2.0",
+  id: 99,
+  method: "session/request_permission",
+  params: {
+    sessionId: "session-1",
+    toolCall: {
+      toolCallId: "permission-tool",
+      title: "Permission tool",
+    },
+    options: [
+      {
+        kind: "allow_once",
+        name: "Allow once",
+        optionId: "allow",
+      },
+    ],
+  },
+} satisfies AnyMessage;
+
+const permissionResponse = {
+  jsonrpc: "2.0",
+  id: 99,
+  result: {
+    outcome: {
+      outcome: "selected",
+      optionId: "allow",
+    },
+  },
+} satisfies AnyMessage;
+
 describe("createHttpStream", () => {
   it("posts initialize with custom headers, opens connection SSE, and emits the initialize response", async () => {
     const controlledFetch = createControlledFetch();
@@ -137,6 +179,68 @@ describe("createHttpStream", () => {
     }
   });
 
+  it("opens session SSE before posting session/load for an existing session", async () => {
+    const controlledFetch = createControlledFetch();
+    const stream = createHttpStream("https://agent.example/acp", {
+      fetch: controlledFetch.fetch,
+    });
+    const writer = stream.writable.getWriter();
+    const reader = stream.readable.getReader();
+
+    try {
+      await writer.write(initializeRequest);
+      await readMessage(reader);
+      await writer.write(loadSessionRequest);
+
+      const sessionGet = requestAt(controlledFetch.requests, 2);
+      const loadPost = requestAt(controlledFetch.requests, 3);
+
+      expect(sessionGet.method).toBe("GET");
+      expect(sessionGet.headers.get(HEADER_CONNECTION_ID)).toBe("connection-1");
+      expect(sessionGet.headers.get(HEADER_SESSION_ID)).toBe(
+        "existing-session",
+      );
+      expect(loadPost.method).toBe("POST");
+      expect(loadPost.headers.get(HEADER_CONNECTION_ID)).toBe("connection-1");
+      expect(loadPost.headers.get(HEADER_SESSION_ID)).toBe("existing-session");
+    } finally {
+      reader.releaseLock();
+      writer.releaseLock();
+      await stream.writable.close();
+    }
+  });
+
+  it("includes the session header on responses to session-scoped server requests", async () => {
+    const controlledFetch = createControlledFetch();
+    const stream = createHttpStream("https://agent.example/acp", {
+      fetch: controlledFetch.fetch,
+    });
+    const writer = stream.writable.getWriter();
+    const reader = stream.readable.getReader();
+
+    try {
+      await writer.write(initializeRequest);
+      await readMessage(reader);
+      await controlledFetch.sendSse(0, sessionNewResponse);
+      await readMessage(reader);
+      await controlledFetch.sendSse(1, permissionRequest);
+      await readMessage(reader);
+      await writer.write(permissionResponse);
+
+      const responsePost = requestAt(controlledFetch.requests, 3);
+      expect(responsePost.method).toBe("POST");
+      expect(responsePost.headers.get(HEADER_CONNECTION_ID)).toBe(
+        "connection-1",
+      );
+      expect(responsePost.headers.get(HEADER_SESSION_ID)).toBe("session-1");
+      expect(JSON.parse(responsePost.body)).toEqual(permissionResponse);
+    } finally {
+      reader.releaseLock();
+      writer.releaseLock();
+      await stream.writable.close();
+    }
+  });
+
   it("propagates cookies across initialize, SSE, session POST, and DELETE", async () => {
     const controlledFetch = createControlledFetch({
       initializeCookies: ["transport=alpha; Path=/"],

src/http-stream.ts

@@ -5,6 +5,7 @@ import {
   HEADER_SESSION_ID,
   JSON_MIME_TYPE,
   isInitializeRequest,
+  messageIdKey,
   sessionIdFromMessageParams,
   sessionIdFromResponseResult,
 } from "./protocol.js";
@@ -44,6 +45,7 @@ class HttpStreamTransport {
   private readonly cookieJar = new ConnectionCookieJar();
   private readonly abortController = new AbortController();
   private readonly knownSessions = new Set<string>();
+  private readonly pendingResponseSessions = new Map<string, string>();
 
   private readableController:
     | ReadableStreamDefaultController<AnyMessage>
@@ -133,7 +135,11 @@ class HttpStreamTransport {
       throw new Error("ACP HTTP stream is not initialized");
     }
 
-    const sessionId = sessionIdFromMessageParams(message);
+    const sessionId = this.sessionIdForOutboundMessage(message);
+    if (sessionId) {
+      this.openSessionSse(sessionId);
+    }
+
     const response = await this.fetchRequest({
       method: "POST",
       headers: {
@@ -149,6 +155,20 @@ class HttpStreamTransport {
     }
   }
 
+  private sessionIdForOutboundMessage(message: AnyMessage): string | undefined {
+    const paramsSessionId = sessionIdFromMessageParams(message);
+    if (paramsSessionId) {
+      return paramsSessionId;
+    }
+
+    if (!("id" in message) || "method" in message) {
+      return undefined;
+    }
+
+    const key = messageIdKey(message.id);
+    return key ? this.pendingResponseSessions.get(key) : undefined;
+  }
+
   private openConnectionSse(): void {
     const connectionId = this.connectionId;
     if (!connectionId) {
@@ -207,6 +227,7 @@ class HttpStreamTransport {
           this.openSessionSse(sessionId);
         }
 
+        this.trackServerRequestRoute(message, headers[HEADER_SESSION_ID]);
         this.enqueue(message);
       }
     } catch (error) {
@@ -218,6 +239,20 @@ class HttpStreamTransport {
     }
   }
 
+  private trackServerRequestRoute(
+    message: AnyMessage,
+    streamSessionId: string | undefined,
+  ): void {
+    if (!streamSessionId || !("method" in message) || !("id" in message)) {
+      return;
+    }
+
+    const key = messageIdKey(message.id);
+    if (key) {
+      this.pendingResponseSessions.set(key, streamSessionId);
+    }
+  }
+
   private async fetchRequest(init: RequestInit): Promise<Response> {
     const response = await this.fetchImpl(this.serverUrl, {
       ...init,

src/server-permission.test.ts

@@ -45,6 +45,64 @@ function createPromptRequest(id: number, sessionId: string) {
 }
 
 describe("AcpServer permission requests over HTTP", () => {
+  it("rejects session-scoped client responses without a session header", async () => {
+    const server = await startTestServer(
+      (conn: AgentSideConnection) =>
+        new TestAgent(conn, { enablePermission: true }),
+    );
+
+    try {
+      const connectionId = await initialize(server.url);
+      const sessionId = await createSession(server.url, connectionId);
+      const sessionSse = await openSessionSse(
+        server.url,
+        connectionId,
+        sessionId,
+      );
+      const sessionEvents = createSseMessageIterator(sessionSse);
+
+      expect(
+        await postJson(server.url, createPromptRequest(3, sessionId), {
+          [HEADER_CONNECTION_ID]: connectionId,
+          [HEADER_SESSION_ID]: sessionId,
+        }),
+      ).toMatchObject({ status: 202 });
+
+      await readNextSseMessage(sessionEvents);
+      const permissionRequest = await readNextSseMessage(sessionEvents);
+
+      const permissionResponse = {
+        jsonrpc: "2.0",
+        id: readMessageId(permissionRequest),
+        result: {
+          outcome: {
+            outcome: "selected",
+            optionId: "allow",
+          },
+        },
+      };
+
+      expect(
+        await postJson(server.url, permissionResponse, {
+          [HEADER_CONNECTION_ID]: connectionId,
+        }),
+      ).toMatchObject({ status: 400 });
+      expect(
+        await postJson(server.url, permissionResponse, {
+          [HEADER_CONNECTION_ID]: connectionId,
+          [HEADER_SESSION_ID]: sessionId,
+        }),
+      ).toMatchObject({ status: 202 });
+
+      await readNextSseMessage(sessionEvents);
+      await readNextSseMessage(sessionEvents);
+      await sessionEvents.return?.();
+      await sessionSse.body?.cancel();
+    } finally {
+      await server.close();
+    }
+  }, 10_000);
+
   it("routes permission requests over session SSE and accepts client responses", async () => {
     const server = await startTestServer(
       (conn: AgentSideConnection) =>
@@ -122,7 +180,10 @@ describe("AcpServer permission requests over HTTP", () => {
               },
             },
           },
-          { [HEADER_CONNECTION_ID]: connectionId },
+          {
+            [HEADER_CONNECTION_ID]: connectionId,
+            [HEADER_SESSION_ID]: sessionId,
+          },
         ),
       ).toMatchObject({ status: 202 });
 

src/server.ts

@@ -255,7 +255,7 @@ export class AcpServer {
     headers: Headers,
   ): Promise<ForwardResult> {
     if (isResponseMessage(message)) {
-      return await forwardClientResponse(connection, message);
+      return await forwardClientResponse(connection, message, headers);
     }
 
     return await forwardClientMethodMessage(connection, message, headers);
@@ -351,7 +351,32 @@ async function forwardClientMethodMessage(
 async function forwardClientResponse(
   connection: ConnectionState,
   message: AnyResponse,
+  headers: Headers,
 ): Promise<ForwardResult> {
+  const key = messageIdKey(message.id);
+  const route = key ? connection.clientResponseRoutes.get(key) : undefined;
+  const headerSessionId = headers.get(HEADER_SESSION_ID);
+
+  if (route && route !== "connection" && !headerSessionId) {
+    return {
+      ok: false,
+      status: 400,
+      message: "Missing Acp-Session-Id",
+    };
+  }
+
+  if (route && route !== "connection" && headerSessionId !== route.session) {
+    return {
+      ok: false,
+      status: 400,
+      message: "Mismatched Acp-Session-Id",
+    };
+  }
+
+  if (key) {
+    connection.clientResponseRoutes.delete(key);
+  }
+
   await writeInbound(connection, message);
   return { ok: true };
 }