Fix: Add Session ID and Header Support to SSE Transport (#660)

.changeset/stupid-plums-wear.md

@@ -0,0 +1,7 @@
+---
+"agents": patch
+---
+
+fix: add session ID and header support to SSE transport
+
+The SSE transport now properly forwards session IDs and request headers to MCP message handlers, achieving closer header parity with StreamableHTTP transport. This allows MCP servers using SSE to access request headers for session management.

packages/agents/src/mcp/index.ts

@@ -1,7 +1,10 @@
 import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
-import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
+import type {
+  JSONRPCMessage,
+  MessageExtraInfo
+} from "@modelcontextprotocol/sdk/types.js";
 import {
   JSONRPCMessageSchema,
   isJSONRPCError,
@@ -85,11 +88,32 @@ export abstract class McpAgent<
     return websockets[0];
   }
 
+  /** Build MessageExtraInfo from extracted headers and auth info */
+  private buildMessageExtraInfo(extraInfo?: {
+    headers: Record<string, string>;
+  }): MessageExtraInfo | undefined {
+    if (!extraInfo) return;
+
+    const headers = { ...extraInfo.headers };
+
+    // Remove internal headers that are not part of the original request
+    delete headers[MCP_HTTP_METHOD_HEADER];
+    delete headers[MCP_MESSAGE_HEADER];
+    delete headers["upgrade"];
+
+    return {
+      requestInfo: { headers }
+    };
+  }
+
   /** Returns a new transport matching the type of the Agent. */
   private initTransport() {
     switch (this.getTransportType()) {
       case "sse": {
-        return new McpSSETransport(() => this.getWebSocket());
+        return new McpSSETransport({
+          getWebSocket: () => this.getWebSocket(),
+          sessionId: this.getSessionId()
+        });
       }
       case "streamable-http": {
         return new StreamableHTTPServerTransport({});
@@ -188,7 +212,8 @@ export abstract class McpAgent<
   /** Handles MCP Messages for the legacy SSE transport. */
   async onSSEMcpMessage(
     _sessionId: string,
-    messageBody: unknown
+    messageBody: unknown,
+    extraInfo?: { headers: Record<string, string> }
   ): Promise<Error | null> {
     // Since we address the DO via both the protocol and the session id,
     // this should never happen, but let's enforce it just in case
@@ -210,7 +235,10 @@ export abstract class McpAgent<
         return null; // Message was handled by elicitation system
       }
 
-      this._transport?.onmessage?.(parsedMessage);
+      // Build extra info with auth and request headers, matching StreamableHTTP behavior
+      const extra = this.buildMessageExtraInfo(extraInfo);
+
+      this._transport?.onmessage?.(parsedMessage, extra);
       return null;
     } catch (error) {
       console.error("Error forwarding message to SSE:", error);

packages/agents/src/mcp/transport.ts

@@ -15,17 +15,23 @@ import type { McpAgent } from ".";
 import { MessageType } from "../ai-types";
 import { MCP_HTTP_METHOD_HEADER, MCP_MESSAGE_HEADER } from "./utils";
 
+export interface McpSSETransportOptions {
+  getWebSocket: () => WebSocket | null;
+  sessionId?: string;
+}
+
 export class McpSSETransport implements Transport {
   sessionId?: string;
   // Set by the server in `server.connect(transport)`
   onclose?: () => void;
   onerror?: (error: Error) => void;
-  onmessage?: (message: JSONRPCMessage) => void;
+  onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void;
 
   private _getWebSocket: () => WebSocket | null;
   private _started = false;
-  constructor(getWebSocket: () => WebSocket | null) {
-    this._getWebSocket = getWebSocket;
+  constructor(options: McpSSETransportOptions) {
+    this._getWebSocket = options.getWebSocket;
+    this.sessionId = options.sessionId;
   }
 
   async start() {

packages/agents/src/mcp/utils.ts

@@ -682,7 +682,9 @@ export const createLegacySseHandler = (
       });
 
       const messageBody = await request.json();
-      const error = await agent.onSSEMcpMessage(sessionId, messageBody);
+      const error = await agent.onSSEMcpMessage(sessionId, messageBody, {
+        headers: Object.fromEntries(request.headers.entries())
+      });
 
       if (error) {
         return new Response(error.message, {

packages/agents/src/tests/mcp/transports/sse.test.ts

@@ -1,5 +1,9 @@
 import { createExecutionContext, env } from "cloudflare:test";
 import { describe, expect, it } from "vitest";
+import type {
+  CallToolResult,
+  JSONRPCResponse
+} from "@modelcontextprotocol/sdk/types.js";
 import worker, { type Env } from "../../worker";
 import { establishSSEConnection } from "../../shared/test-utils";
 
@@ -117,4 +121,77 @@ describe("SSE Transport", () => {
       expect(response.headers.get("Content-Type")).toBe("text/event-stream");
     });
   });
+
+  describe("Header and Auth Handling", () => {
+    it("should pass headers and session ID to transport via requestInfo", async () => {
+      const ctx = createExecutionContext();
+      const { sessionId, reader } = await establishSSEConnection(ctx);
+
+      // Send request with custom headers using the echoRequestInfo tool
+      const request = new Request(`${baseUrl}/message?sessionId=${sessionId}`, {
+        method: "POST",
+        body: JSON.stringify({
+          id: "echo-headers-1",
+          jsonrpc: "2.0",
+          method: "tools/call",
+          params: {
+            name: "echoRequestInfo",
+            arguments: {}
+          }
+        }),
+        headers: {
+          "Content-Type": "application/json",
+          "x-user-id": "test-user-123",
+          "x-request-id": "req-456",
+          "x-custom-header": "custom-value"
+        }
+      });
+
+      const response = await worker.fetch(request, env, ctx);
+      expect(response.status).toBe(202); // SSE returns 202 Accepted
+
+      // Read the response from the SSE stream
+      const { value } = await reader.read();
+      const event = new TextDecoder().decode(value);
+      const lines = event.split("\n");
+      expect(lines[0]).toEqual("event: message");
+
+      // Parse the JSON response from the data line
+      const dataLine = lines.find((line) => line.startsWith("data:"));
+      const parsed = JSON.parse(
+        dataLine!.replace("data: ", "")
+      ) as JSONRPCResponse;
+      expect(parsed.id).toBe("echo-headers-1");
+
+      // Extract the echoed request info
+      const result = parsed.result as CallToolResult;
+      const textContent = result.content?.[0];
+      if (!textContent || textContent.type !== "text") {
+        throw new Error("Expected text content in tool result");
+      }
+      const echoedData = JSON.parse(textContent.text);
+
+      // Verify custom headers were passed through
+      expect(echoedData.hasRequestInfo).toBe(true);
+      expect(echoedData.headers["x-user-id"]).toBe("test-user-123");
+      expect(echoedData.headers["x-request-id"]).toBe("req-456");
+      expect(echoedData.headers["x-custom-header"]).toBe("custom-value");
+
+      // Verify that certain internal headers that the transport adds are NOT exposed
+      // The transport filters cf-mcp-method, cf-mcp-message, and upgrade headers
+      expect(echoedData.headers["cf-mcp-method"]).toBeUndefined();
+      expect(echoedData.headers["cf-mcp-message"]).toBeUndefined();
+      expect(echoedData.headers["upgrade"]).toBeUndefined();
+
+      // Verify standard headers are also present
+      expect(echoedData.headers["content-type"]).toBe("application/json");
+
+      // Check what properties are available in extra
+      expect(echoedData.availableExtraKeys).toBeDefined();
+
+      // Verify sessionId is passed through extra data
+      expect(echoedData.sessionId).toBeDefined();
+      expect(echoedData.sessionId).toBe(sessionId);
+    });
+  });
 });

packages/agents/src/tests/mcp/transports/streamable-http.test.ts

@@ -583,4 +583,67 @@
       expect(tools.some((t) => t.name === "temp-echo")).toBe(false);
     });
   });
+
+  describe("Header and Auth Handling", () => {
+    it("should pass custom headers to transport via requestInfo", async () => {
+      const ctx = createExecutionContext();
+      const sessionId = await initializeStreamableHTTPServer(ctx);
+
+      // Send request with custom headers using the echoRequestInfo tool
+      const echoMessage: JSONRPCMessage = {
+        id: "echo-headers-1",
+        jsonrpc: "2.0",
+        method: "tools/call",
+        params: {
+          name: "echoRequestInfo",
+          arguments: {}
+        }
+      };
+
+      const request = new Request(baseUrl, {
+        body: JSON.stringify(echoMessage),
+        headers: {
+          Accept: "application/json, text/event-stream",
+          "Content-Type": "application/json",
+          "mcp-session-id": sessionId,
+          "x-user-id": "test-user-123",
+          "x-request-id": "req-456",
+          "x-custom-header": "custom-value"
+        },
+        method: "POST"
+      });
+
+      const response = await worker.fetch(request, env, ctx);
+      expect(response.status).toBe(200);
+
+      // Parse the SSE response
+      const sseText = await readSSEEvent(response);
+      const parsed = parseSSEData(sseText) as JSONRPCResponse;
+      expect(parsed.id).toBe("echo-headers-1");
+
+      // Extract the echoed request info
+      const result = parsed.result as CallToolResult;
+      const echoedData = JSON.parse(result.content?.[0]?.text || "{}");
+
+      // Verify custom headers were passed through
+      expect(echoedData.hasRequestInfo).toBe(true);
+      expect(echoedData.headers["x-user-id"]).toBe("test-user-123");
+      expect(echoedData.headers["x-request-id"]).toBe("req-456");
+      expect(echoedData.headers["x-custom-header"]).toBe("custom-value");
+
+      // Verify that certain internal headers that the transport adds are NOT exposed
+      // The transport adds cf-mcp-method and cf-mcp-message internally but should filter them
+      expect(echoedData.headers["cf-mcp-method"]).toBeUndefined();
+      expect(echoedData.headers["cf-mcp-message"]).toBeUndefined();
+      expect(echoedData.headers["upgrade"]).toBeUndefined();
+
+      // Verify standard headers are also present
+      expect(echoedData.headers["accept"]).toContain("text/event-stream");
+      expect(echoedData.headers["content-type"]).toBe("application/json");
+
+      // Verify sessionId is passed through extra data
+      expect(echoedData.sessionId).toBeDefined();
+      expect(echoedData.sessionId).toBe(sessionId);
+    });
+  });
 });

packages/agents/src/tests/worker.ts

@@ -169,6 +169,56 @@ export class TestMcpAgent extends McpAgent<Env, State, Props> {
         return { content: [{ type: "text", text: "nothing to remove" }] };
       }
     );
+
+    // Echo request info for testing header and auth passthrough
+    this.server.tool(
+      "echoRequestInfo",
+      "Echo back request headers and auth info",
+      {},
+      async (_args, extra): Promise<CallToolResult> => {
+        // Extract headers from requestInfo, auth from authInfo
+        const headers = extra?.requestInfo?.headers || {};
+        const authInfo = extra?.authInfo || null;
+
+        // Get all properties from extra, excluding functions
+        const extraKeys = extra
+          ? Object.keys(extra).filter(
+              (key) => typeof extra[key as keyof typeof extra] !== "function"
+            )
+          : [];
+
+        // Build response object with all available data
+        const responseData: Record<string, any> = {
+          headers: headers,
+          authInfo: authInfo,
+          hasRequestInfo: !!extra?.requestInfo,
+          hasAuthInfo: !!extra?.authInfo,
+          requestId: extra?.requestId,
+          // Include any sessionId if it exists
+          sessionId: (extra as any)?.sessionId || null,
+          // List all available properties in extra
+          availableExtraKeys: extraKeys
+        };
+
+        // Add any other properties from extra that aren't already included
+        extraKeys.forEach((key) => {
+          if (
+            !["requestInfo", "authInfo", "requestId", "sessionId"].includes(key)
+          ) {
+            responseData[`extra_${key}`] = (extra as any)[key];
+          }
+        });
+
+        return {
+          content: [
+            {
+              type: "text",
+              text: JSON.stringify(responseData, null, 2)
+            }
+          ]
+        };
+      }
+    );
   }
 }