From e3bb99cd5c9f25e079f678b3d82c05765d02e12f Mon Sep 17 00:00:00 2001 From: ihrpr Date: Sun, 6 Apr 2025 17:27:00 +0100 Subject: [PATCH] remove the use of enableSessionManagement --- src/server/streamable-http.test.ts | 141 ++++++++++++----------------- src/server/streamable-http.ts | 90 +++++++++--------- 2 files changed, 104 insertions(+), 127 deletions(-) diff --git a/src/server/streamable-http.test.ts b/src/server/streamable-http.test.ts index 74b92c1a..b888b0ec 100644 --- a/src/server/streamable-http.test.ts +++ b/src/server/streamable-http.test.ts @@ -2,7 +2,7 @@ import { IncomingMessage, ServerResponse } from "node:http"; import { StreamableHTTPServerTransport } from "./streamable-http.js"; import { JSONRPCMessage } from "../types.js"; import { Readable } from "node:stream"; - +import { randomUUID } from "node:crypto"; // Mock IncomingMessage function createMockRequest(options: { method: string; @@ -10,7 +10,7 @@ function createMockRequest(options: { body?: string; }): IncomingMessage { const readable = new Readable(); - readable._read = () => {}; + readable._read = () => { }; if (options.body) { readable.push(options.body); readable.push(null); @@ -37,12 +37,13 @@ function createMockResponse(): jest.Mocked { } describe("StreamableHTTPServerTransport", () => { - const endpoint = "/mcp"; let transport: StreamableHTTPServerTransport; let mockResponse: jest.Mocked; beforeEach(() => { - transport = new StreamableHTTPServerTransport(endpoint); + transport = new StreamableHTTPServerTransport({ + sessionId: randomUUID(), + }); mockResponse = createMockResponse(); }); @@ -61,7 +62,7 @@ describe("StreamableHTTPServerTransport", () => { const initializeMessage: JSONRPCMessage = { jsonrpc: "2.0", method: "initialize", - params: { + params: { clientInfo: { name: "test-client", version: "1.0" }, protocolVersion: "2025-03-26" }, @@ -119,49 +120,13 @@ describe("StreamableHTTPServerTransport", () => { expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"jsonrpc":"2.0"')); expect(mockResponse.end).toHaveBeenCalledWith(expect.stringContaining('"message":"Bad Request: Mcp-Session-Id header is required"')); }); - - it("should always include session ID in initialization response even in stateless mode", async () => { - // Create a stateless transport for this test - const statelessTransport = new StreamableHTTPServerTransport(endpoint, { enableSessionManagement: false }); - - // Create an initialization request - const initializeMessage: JSONRPCMessage = { - jsonrpc: "2.0", - method: "initialize", - params: { - clientInfo: { name: "test-client", version: "1.0" }, - protocolVersion: "2025-03-26" - }, - id: "init-1", - }; - - const req = createMockRequest({ - method: "POST", - headers: { - "content-type": "application/json", - "accept": "application/json", - }, - body: JSON.stringify(initializeMessage), - }); - - await statelessTransport.handleRequest(req, mockResponse); - - // In stateless mode, session ID should also be included for initialize responses - expect(mockResponse.writeHead).toHaveBeenCalledWith( - 200, - expect.objectContaining({ - "mcp-session-id": statelessTransport.sessionId, - }) - ); - }); }); - describe("Stateless Mode", () => { let statelessTransport: StreamableHTTPServerTransport; let mockResponse: jest.Mocked; beforeEach(() => { - statelessTransport = new StreamableHTTPServerTransport(endpoint, { enableSessionManagement: false }); + statelessTransport = new StreamableHTTPServerTransport({ sessionId: undefined }); mockResponse = createMockResponse(); }); @@ -268,7 +233,7 @@ describe("StreamableHTTPServerTransport", () => { }); await statelessTransport.handleRequest(req2, mockResponse); - + // Should still succeed expect(mockResponse.writeHead).toHaveBeenCalledWith( 200, @@ -278,12 +243,12 @@ describe("StreamableHTTPServerTransport", () => { ); }); - it("should handle initialization requests properly in both modes", async () => { + it("should handle initialization requests properly in statefull mode", async () => { // Initialize message that would typically be sent during initialization const initializeMessage: JSONRPCMessage = { jsonrpc: "2.0", method: "initialize", - params: { + params: { clientInfo: { name: "test-client", version: "1.0" }, protocolVersion: "2025-03-26" }, @@ -301,7 +266,7 @@ describe("StreamableHTTPServerTransport", () => { }); await transport.handleRequest(statefulReq, mockResponse); - + // In stateful mode, session ID should be included in the response header expect(mockResponse.writeHead).toHaveBeenCalledWith( 200, @@ -309,9 +274,19 @@ describe("StreamableHTTPServerTransport", () => { "mcp-session-id": transport.sessionId, }) ); + }); - // Reset mocks for stateless test - mockResponse.writeHead.mockClear(); + it("should handle initialization requests properly in stateless mode", async () => { + // Initialize message that would typically be sent during initialization + const initializeMessage: JSONRPCMessage = { + jsonrpc: "2.0", + method: "initialize", + params: { + clientInfo: { name: "test-client", version: "1.0" }, + protocolVersion: "2025-03-26" + }, + id: "init-1", + }; // Test stateless transport const statelessReq = createMockRequest({ @@ -324,10 +299,11 @@ describe("StreamableHTTPServerTransport", () => { }); await statelessTransport.handleRequest(statelessReq, mockResponse); - + // In stateless mode, session ID should also be included for initialize responses const headers = mockResponse.writeHead.mock.calls[0][1]; - expect(headers).toHaveProperty("mcp-session-id", statelessTransport.sessionId); + expect(headers).not.toHaveProperty("mcp-session-id"); + }); }); @@ -519,14 +495,14 @@ describe("StreamableHTTPServerTransport", () => { // Send a message to first connection const message1: JSONRPCMessage = { - jsonrpc: "2.0", - method: "test1", - params: {}, + jsonrpc: "2.0", + method: "test1", + params: {}, id: 1 }; - + await transport.send(message1); - + // Get message ID (captured from write call) const writeCall = mockResponse.write.mock.calls[0][0] as string; const idMatch = writeCall.match(/id: ([a-f0-9-]+)/); @@ -550,12 +526,12 @@ describe("StreamableHTTPServerTransport", () => { // Send a second message const message2: JSONRPCMessage = { - jsonrpc: "2.0", - method: "test2", - params: {}, + jsonrpc: "2.0", + method: "test2", + params: {}, id: 2 }; - + await transport.send(message2); // Verify the second message was received by both connections @@ -596,7 +572,7 @@ describe("StreamableHTTPServerTransport", () => { params: {}, id: "test-id", }; - + const reqPost = createMockRequest({ method: "POST", headers: { @@ -605,28 +581,28 @@ describe("StreamableHTTPServerTransport", () => { }, body: JSON.stringify(requestMessage), }); - + await transport.handleRequest(reqPost, mockResponse1); - + // Send a response with matching ID const responseMessage: JSONRPCMessage = { jsonrpc: "2.0", result: { success: true }, id: "test-id", }; - + await transport.send(responseMessage); - + // Verify response was sent to the right connection expect(mockResponse1.write).toHaveBeenCalledWith( expect.stringContaining(JSON.stringify(responseMessage)) ); - + // Check if write was called with this exact message on the second connection - const writeCallsOnSecondConn = mockResponse2.write.mock.calls.filter(call => + const writeCallsOnSecondConn = mockResponse2.write.mock.calls.filter(call => typeof call[0] === 'string' && call[0].includes(JSON.stringify(responseMessage)) ); - + // Verify the response wasn't broadcast to all connections expect(writeCallsOnSecondConn.length).toBe(0); }); @@ -680,7 +656,7 @@ describe("StreamableHTTPServerTransport", () => { const message: JSONRPCMessage = { jsonrpc: "2.0", method: "initialize", - params: { + params: { clientInfo: { name: "test-client", version: "1.0" }, protocolVersion: "2025-03-26" }, @@ -715,17 +691,17 @@ describe("StreamableHTTPServerTransport", () => { it("should handle pre-parsed batch messages", async () => { const batchMessages: JSONRPCMessage[] = [ - { - jsonrpc: "2.0", - method: "method1", + { + jsonrpc: "2.0", + method: "method1", params: { data: "test1" }, - id: "batch1" + id: "batch1" }, - { - jsonrpc: "2.0", - method: "method2", + { + jsonrpc: "2.0", + method: "method2", params: { data: "test2" }, - id: "batch2" + id: "batch2" }, ]; @@ -800,7 +776,7 @@ describe("StreamableHTTPServerTransport", () => { let mockResponse: jest.Mocked; beforeEach(() => { - transportWithHeaders = new StreamableHTTPServerTransport(endpoint, { customHeaders }); + transportWithHeaders = new StreamableHTTPServerTransport({ sessionId: randomUUID(), customHeaders }); mockResponse = createMockResponse(); }); @@ -875,9 +851,10 @@ describe("StreamableHTTPServerTransport", () => { }); it("should not override essential headers with custom headers", async () => { - const transportWithConflictingHeaders = new StreamableHTTPServerTransport(endpoint, { + const transportWithConflictingHeaders = new StreamableHTTPServerTransport({ + sessionId: randomUUID(), customHeaders: { - "Content-Type": "text/plain", // 尝试覆盖必要的 Content-Type 头 + "Content-Type": "text/plain", "X-Custom-Header": "custom-value" } }); @@ -902,8 +879,10 @@ describe("StreamableHTTPServerTransport", () => { }); it("should work with empty custom headers", async () => { - const transportWithoutHeaders = new StreamableHTTPServerTransport(endpoint); - + const transportWithoutHeaders = new StreamableHTTPServerTransport({ + sessionId: randomUUID(), + }); + const req = createMockRequest({ method: "GET", headers: { diff --git a/src/server/streamable-http.ts b/src/server/streamable-http.ts index 13d28968..ef654f85 100644 --- a/src/server/streamable-http.ts +++ b/src/server/streamable-http.ts @@ -23,11 +23,11 @@ interface StreamConnection { */ export interface StreamableHTTPServerTransportOptions { /** - * Whether to enable session management through mcp-session-id headers - * When set to false, the transport operates in stateless mode without session validation - * @default true + * The session ID SHOULD be globally unique and cryptographically secure (e.g., a securely generated UUID, a JWT, or a cryptographic hash) + * + * When sessionId is not set, the transport will be in stateless mode. */ - enableSessionManagement?: boolean; + sessionId: string | undefined; /** * Custom headers to be included in all responses @@ -38,17 +38,19 @@ export interface StreamableHTTPServerTransportOptions { /** * Server transport for Streamable HTTP: this implements the MCP Streamable HTTP transport specification. - * It supports both SSE streaming and direct HTTP responses, with session management and message resumability. + * It supports both SSE streaming and direct HTTP responses. * * Usage example: * * ```typescript - * // Stateful mode (default) - with session management - * const statefulTransport = new StreamableHTTPServerTransport("/mcp"); + * // Stateful mode - server sets the session ID + * const statefulTransport = new StreamableHTTPServerTransport({ + * sessionId: randomUUID(), + * }); * - * // Stateless mode - without session management - * const statelessTransport = new StreamableHTTPServerTransport("/mcp", { - * enableSessionManagement: false + * // Stateless mode - explisitly set session ID to undefined + * const statelessTransport = new StreamableHTTPServerTransport({ + * sessionId: undefined, * }); * * // Using with pre-parsed request body @@ -70,23 +72,22 @@ export interface StreamableHTTPServerTransportOptions { */ export class StreamableHTTPServerTransport implements Transport { private _connections: Map = new Map(); - private _sessionId: string; + // when sessionID is not set, it means the transport is in stateless mode + private _sessionId: string | undefined; private _messageHistory: Map = new Map(); private _started: boolean = false; private _requestConnections: Map = new Map(); // request ID to connection ID mapping - private _enableSessionManagement: boolean; private _customHeaders: Record; onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage) => void; - constructor(private _endpoint: string, options?: StreamableHTTPServerTransportOptions) { - this._sessionId = randomUUID(); - this._enableSessionManagement = options?.enableSessionManagement !== false; + constructor(options: StreamableHTTPServerTransportOptions) { + this._sessionId = options?.sessionId; this._customHeaders = options?.customHeaders || {}; } @@ -106,13 +107,13 @@ export class StreamableHTTPServerTransport implements Transport { */ async handleRequest(req: IncomingMessage, res: ServerResponse, parsedBody?: unknown): Promise { // Only validate session ID for non-initialization requests when session management is enabled - if (this._enableSessionManagement) { + if (this._sessionId !== undefined) { const sessionId = req.headers["mcp-session-id"]; - + // Check if this might be an initialization request - const isInitializationRequest = req.method === "POST" && + const isInitializationRequest = req.method === "POST" && req.headers["content-type"]?.includes("application/json"); - + if (isInitializationRequest) { // For POST requests with JSON content, we need to check if it's an initialization request // This will be done in handlePostRequest, as we need to parse the body @@ -189,9 +190,9 @@ export class StreamableHTTPServerTransport implements Transport { "Cache-Control": "no-cache", Connection: "keep-alive", }; - - // Only include session ID header if session management is enabled - if (this._enableSessionManagement) { + + // Only include session ID header if session management is enabled by assigning a value to _sessionId + if (this._sessionId !== undefined) { headers["mcp-session-id"] = this._sessionId; } @@ -231,8 +232,8 @@ export class StreamableHTTPServerTransport implements Transport { try { // validate the Accept header const acceptHeader = req.headers.accept; - if (!acceptHeader || - (!acceptHeader.includes("application/json") && !acceptHeader.includes("text/event-stream"))) { + if (!acceptHeader || + (!acceptHeader.includes("application/json") && !acceptHeader.includes("text/event-stream"))) { res.writeHead(406).end(JSON.stringify({ jsonrpc: "2.0", error: { @@ -270,7 +271,7 @@ export class StreamableHTTPServerTransport implements Transport { } let messages: JSONRPCMessage[]; - + // handle batch and single messages if (Array.isArray(rawMessage)) { messages = rawMessage.map(msg => JSONRPCMessageSchema.parse(msg)); @@ -286,7 +287,7 @@ export class StreamableHTTPServerTransport implements Transport { // check if it contains requests const hasRequests = messages.some(msg => 'method' in msg && 'id' in msg); - const hasOnlyNotificationsOrResponses = messages.every(msg => + const hasOnlyNotificationsOrResponses = messages.every(msg => ('method' in msg && !('id' in msg)) || ('result' in msg || 'error' in msg)); if (hasOnlyNotificationsOrResponses) { @@ -300,7 +301,7 @@ export class StreamableHTTPServerTransport implements Transport { } else if (hasRequests) { // if it contains requests, you can choose to return an SSE stream or a JSON response const useSSE = acceptHeader.includes("text/event-stream"); - + if (useSSE) { const headers: Record = { ...this._customHeaders, @@ -308,10 +309,8 @@ export class StreamableHTTPServerTransport implements Transport { "Cache-Control": "no-cache", Connection: "keep-alive", }; - - // Only include session ID header if session management is enabled - // Always include session ID for initialization requests - if (this._enableSessionManagement || isInitializationRequest) { + + if (this._sessionId !== undefined) { headers["mcp-session-id"] = this._sessionId; } @@ -351,20 +350,19 @@ export class StreamableHTTPServerTransport implements Transport { ...this._customHeaders, "Content-Type": "application/json", }; - - // Only include session ID header if session management is enabled - // Always include session ID for initialization requests - if (this._enableSessionManagement || isInitializationRequest) { + + + if (this._sessionId !== undefined) { headers["mcp-session-id"] = this._sessionId; } res.writeHead(200, headers); - + // handle each message for (const message of messages) { this.onmessage?.(message); } - + res.end(); } } @@ -396,11 +394,11 @@ export class StreamableHTTPServerTransport implements Transport { */ private replayMessages(connectionId: string, lastEventId: string): void { if (!lastEventId) return; - + // only replay messages that should be sent on this connection const messages = Array.from(this._messageHistory.entries()) - .filter(([id, { connectionId: msgConnId }]) => - id > lastEventId && + .filter(([id, { connectionId: msgConnId }]) => + id > lastEventId && (!msgConnId || msgConnId === connectionId)) // only replay messages that are not specified to a connection or specified to the current connection .sort(([a], [b]) => a.localeCompare(b)); @@ -430,11 +428,11 @@ export class StreamableHTTPServerTransport implements Transport { } let targetConnectionId = ""; - + // if it is a response, find the corresponding request connection if ('id' in message && ('result' in message || 'error' in message)) { const connId = this._requestConnections.get(String(message.id)); - + // if the corresponding connection is not found, the connection may be disconnected if (!connId || !this._connections.has(connId)) { // select an available connection @@ -458,9 +456,9 @@ export class StreamableHTTPServerTransport implements Transport { } const messageId = randomUUID(); - this._messageHistory.set(messageId, { - message, - connectionId: targetConnectionId + this._messageHistory.set(messageId, { + message, + connectionId: targetConnectionId }); // keep the message history in a reasonable range @@ -490,7 +488,7 @@ export class StreamableHTTPServerTransport implements Transport { /** * Returns the session ID for this transport */ - get sessionId(): string { + get sessionId(): string | undefined { return this._sessionId; } } \ No newline at end of file