From 585dc06708a80920634d2ae7ca76f067c6024874 Mon Sep 17 00:00:00 2001 From: ihrpr Date: Wed, 16 Apr 2025 11:09:47 +0100 Subject: [PATCH] introduce onsessioninitialized to avoid race conditions --- src/client/streamableHttp.ts | 2 +- .../server/jsonResponseStreamableHttp.ts | 13 ++++----- src/examples/server/simpleStreamableHttp.ts | 12 ++++---- .../standaloneSseWithGetStreamableHttp.ts | 12 ++++---- src/server/streamableHttp.test.ts | 29 ++++--------------- src/server/streamableHttp.ts | 23 ++++++++++++--- 6 files changed, 45 insertions(+), 46 deletions(-) diff --git a/src/client/streamableHttp.ts b/src/client/streamableHttp.ts index 0510fe84..9efd8536 100644 --- a/src/client/streamableHttp.ts +++ b/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"; diff --git a/src/examples/server/jsonResponseStreamableHttp.ts b/src/examples/server/jsonResponseStreamableHttp.ts index 9be63f52..101a581f 100644 --- a/src/examples/server/jsonResponseStreamableHttp.ts +++ b/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 diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts index eb28bf4c..a318b738 100644 --- a/src/examples/server/simpleStreamableHttp.ts +++ b/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 diff --git a/src/examples/server/standaloneSseWithGetStreamableHttp.ts b/src/examples/server/standaloneSseWithGetStreamableHttp.ts index 5f96ca82..8c8c3baa 100644 --- a/src/examples/server/standaloneSseWithGetStreamableHttp.ts +++ b/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 diff --git a/src/server/streamableHttp.test.ts b/src/server/streamableHttp.test.ts index efd5de1c..0794e4bb 100644 --- a/src/server/streamableHttp.test.ts +++ b/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, { diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts index 31aad09c..92bbed37 100644 --- a/src/server/streamableHttp.ts +++ b/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