introduce onsessioninitialized to avoid race conditions

Author: ihrpr

src/client/streamableHttp.ts

@@ -1,5 +1,5 @@
 import { Transport } from "../shared/transport.js";
-import { isInitializedNotification, isJSONRPCNotification, isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
+import { isInitializedNotification, isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema } from "../types.js";
 import { auth, AuthResult, OAuthClientProvider, UnauthorizedError } from "./auth.js";
 import { EventSourceParserStream } from "eventsource-parser/stream";
 

src/examples/server/jsonResponseStreamableHttp.ts

@@ -95,18 +95,17 @@ app.post('/mcp', async (req: Request, res: Response) => {
       transport = new StreamableHTTPServerTransport({
         sessionIdGenerator: () => randomUUID(),
         enableJsonResponse: true, // Enable JSON response mode
+        onsessioninitialized: (sessionId) => {
+          // Store the transport by session ID when session is initialized
+          // This avoids race conditions where requests might come in before the session is stored
+          console.log(`Session initialized with ID: ${sessionId}`);
+          transports[sessionId] = transport;
+        }
       });
 
       // Connect the transport to the MCP server BEFORE handling the request
       await server.connect(transport);
-
-      // After handling the request, if we get a session ID back, store the transport
       await transport.handleRequest(req, res, req.body);
-
-      // Store the transport by session ID for future requests
-      if (transport.sessionId) {
-        transports[transport.sessionId] = transport;
-      }
       return; // Already handled
     } else {
       // Invalid request - no session ID or not initialization request

src/examples/server/simpleStreamableHttp.ts

@@ -247,19 +247,19 @@ app.post('/mcp', async (req: Request, res: Response) => {
       transport = new StreamableHTTPServerTransport({
         sessionIdGenerator: () => randomUUID(),
         eventStore, // Enable resumability
+        onsessioninitialized: (sessionId) => {
+          // Store the transport by session ID when session is initialized
+          // This avoids race conditions where requests might come in before the session is stored
+          console.log(`Session initialized with ID: ${sessionId}`);
+          transports[sessionId] = transport;
+        }
       });
 
       // Connect the transport to the MCP server BEFORE handling the request
       // so responses can flow back through the same transport
       await server.connect(transport);
 
-      // After handling the request, if we get a session ID back, store the transport
       await transport.handleRequest(req, res, req.body);
-
-      // Store the transport by session ID for future requests
-      if (transport.sessionId) {
-        transports[transport.sessionId] = transport;
-      }
       return; // Already handled
     } else {
       // Invalid request - no session ID or not initialization request

src/examples/server/standaloneSseWithGetStreamableHttp.ts

@@ -52,17 +52,19 @@ app.post('/mcp', async (req: Request, res: Response) => {
       // New initialization request
       transport = new StreamableHTTPServerTransport({
         sessionIdGenerator: () => randomUUID(),
+        onsessioninitialized: (sessionId) => {
+          // Store the transport by session ID when session is initialized
+          // This avoids race conditions where requests might come in before the session is stored
+          console.log(`Session initialized with ID: ${sessionId}`);
+          transports[sessionId] = transport;
+        }
       });
 
       // Connect the transport to the MCP server
       await server.connect(transport);
 
+      // Handle the request - the onsessioninitialized callback will store the transport
       await transport.handleRequest(req, res, req.body);
-
-      // Store the transport by session ID for future requests
-      if (transport.sessionId) {
-        transports[transport.sessionId] = transport;
-      }
       return; // Already handled
     } else {
       // Invalid request - no session ID or not initialization request

src/server/streamableHttp.test.ts

@@ -187,15 +187,11 @@ describe("StreamableHTTPServerTransport", () => {
     expect(sessionId).toBeDefined();
 
     // Try second initialize
-    const secondInitMessage: JSONRPCMessage = {
-      jsonrpc: "2.0",
-      method: "initialize",
-      params: {
-        clientInfo: { name: "test-client-2", version: "1.0" },
-        protocolVersion: "2025-03-26",
-      },
-      id: "init-2",
+    const secondInitMessage = {
+      ...TEST_MESSAGES.initialize,
+      id: "second-init"
     };
+
     const response = await sendPostRequest(baseUrl, secondInitMessage);
 
     expect(response.status).toBe(400);
@@ -1092,14 +1088,7 @@ describe("StreamableHTTPServerTransport in stateless mode", () => {
   });
 
   it("should handle POST requests with various session IDs in stateless mode", async () => {
-    // Initialize the server first
-    await fetch(baseUrl, {
-      method: "POST",
-      headers: { "Content-Type": "application/json", Accept: "application/json, text/event-stream" },
-      body: JSON.stringify({
-        jsonrpc: "2.0", method: "initialize", params: { clientInfo: { name: "test-client", version: "1.0" }, protocolVersion: "2025-03-26" }, id: "init-1"
-      }),
-    });
+    await sendPostRequest(baseUrl, TEST_MESSAGES.initialize);
 
     // Try with a random session ID - should be accepted
     const response1 = await fetch(baseUrl, {
@@ -1131,13 +1120,7 @@ describe("StreamableHTTPServerTransport in stateless mode", () => {
     // one standalone SSE stream at a time
 
     // Initialize the server first
-    await fetch(baseUrl, {
-      method: "POST",
-      headers: { "Content-Type": "application/json", Accept: "application/json, text/event-stream" },
-      body: JSON.stringify({
-        jsonrpc: "2.0", method: "initialize", params: { clientInfo: { name: "test-client", version: "1.0" }, protocolVersion: "2025-03-26" }, id: "init-1"
-      }),
-    });
+    await sendPostRequest(baseUrl, TEST_MESSAGES.initialize);
 
     // Open first SSE stream
     const stream1 = await fetch(baseUrl, {

src/server/streamableHttp.ts

@@ -1,6 +1,6 @@
 import { IncomingMessage, ServerResponse } from "node:http";
 import { Transport } from "../shared/transport.js";
-import { isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js";
+import { isInitializeRequest, isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js";
 import getRawBody from "raw-body";
 import contentType from "content-type";
 import { randomUUID } from "node:crypto";
@@ -39,6 +39,15 @@ export interface StreamableHTTPServerTransportOptions {
    */
   sessionIdGenerator: () => string | undefined;
 
+  /**
+   * A callback for session initialization events
+   * This is called when the server initializes a new session.
+   * Usefult in cases when you need to register multiple mcp sessions
+   * and need to keep track of them.
+   * @param sessionId The generated session ID
+   */
+  onsessioninitialized?: (sessionId: string) => void;
+
   /**
    * If true, the server will return JSON responses instead of starting an SSE stream.
    * This can be useful for simple request/response scenarios without streaming.
@@ -98,6 +107,7 @@ export class StreamableHTTPServerTransport implements Transport {
   private _enableJsonResponse: boolean = false;
   private _standaloneSseStreamId: string = '_GET_stream';
   private _eventStore?: EventStore;
+  private _onsessioninitialized?: (sessionId: string) => void;
 
   sessionId?: string | undefined;
   onclose?: () => void;
@@ -108,6 +118,7 @@ export class StreamableHTTPServerTransport implements Transport {
     this.sessionIdGenerator = options.sessionIdGenerator;
     this._enableJsonResponse = options.enableJsonResponse ?? false;
     this._eventStore = options.eventStore;
+    this._onsessioninitialized = options.onsessioninitialized;
   }
 
   /**
@@ -328,9 +339,7 @@ export class StreamableHTTPServerTransport implements Transport {
 
       // Check if this is an initialization request
       // https://spec.modelcontextprotocol.io/specification/2025-03-26/basic/lifecycle/
-      const isInitializationRequest = messages.some(
-        msg => 'method' in msg && msg.method === 'initialize'
-      );
+      const isInitializationRequest = messages.some(isInitializeRequest);
       if (isInitializationRequest) {
         // If it's a server with session management and the session ID is already set we should reject the request
         // to avoid re-initialization.
@@ -359,6 +368,12 @@ export class StreamableHTTPServerTransport implements Transport {
         this.sessionId = this.sessionIdGenerator();
         this._initialized = true;
 
+        // If we have a session ID and an onsessioninitialized handler, call it immediately
+        // This is needed in cases where the server needs to keep track of multiple sessions
+        if (this.sessionId && this._onsessioninitialized) {
+          this._onsessioninitialized(this.sessionId);
+        }
+
       }
       // If an Mcp-Session-Id is returned by the server during initialization,
       // clients using the Streamable HTTP transport MUST include it