From 44c9bffa8511372b760dc5edb78650563053b18f Mon Sep 17 00:00:00 2001 From: Jason Victor Date: Sun, 23 Feb 2025 12:39:15 -0500 Subject: [PATCH 1/4] added session id to request handler extra data --- src/server/sse.ts | 2 -- src/shared/protocol.ts | 13 ++++++++++++- src/shared/transport.ts | 2 ++ 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/src/server/sse.ts b/src/server/sse.ts index 84c1cbb9..1bc304a8 100644 --- a/src/server/sse.ts +++ b/src/server/sse.ts @@ -136,8 +136,6 @@ export class SSEServerTransport implements Transport { /** * Returns the session ID for this transport. - * - * This can be used to route incoming POST requests. */ get sessionId(): string { return this._sessionId; diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index a4f211c6..ef303a39 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -74,6 +74,11 @@ export type RequestHandlerExtra = { * An abort signal used to communicate if the request was cancelled from the sender's side. */ signal: AbortSignal; + + /** + * The session ID from the transport, if available. + */ + sessionId?: string; }; /** @@ -239,9 +244,15 @@ export abstract class Protocol< const abortController = new AbortController(); this._requestHandlerAbortControllers.set(request.id, abortController); + // Create extra object with both abort signal and sessionId from transport + const extra: RequestHandlerExtra = { + signal: abortController.signal, + sessionId: this._transport?.sessionId, + }; + // Starting with Promise.resolve() puts any synchronous errors into the monad as well. Promise.resolve() - .then(() => handler(request, { signal: abortController.signal })) + .then(() => handler(request, extra)) .then( (result) => { if (abortController.signal.aborted) { diff --git a/src/shared/transport.ts b/src/shared/transport.ts index 5843cf00..bcadcf4d 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -41,4 +41,6 @@ export interface Transport { * Callback for when a message (request or response) is received over the connection. */ onmessage?: (message: JSONRPCMessage) => void; + + sessionId?: string; } From 0408fdbb6284d4609703fcbc374082f653200ae7 Mon Sep 17 00:00:00 2001 From: Jason Victor Date: Mon, 24 Feb 2025 11:06:34 -0500 Subject: [PATCH 2/4] put missing comment back --- src/server/sse.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/server/sse.ts b/src/server/sse.ts index 1bc304a8..84c1cbb9 100644 --- a/src/server/sse.ts +++ b/src/server/sse.ts @@ -136,6 +136,8 @@ export class SSEServerTransport implements Transport { /** * Returns the session ID for this transport. + * + * This can be used to route incoming POST requests. */ get sessionId(): string { return this._sessionId; From 57f3b07bb9a827ab97089bcbeeca89f56b0c009c Mon Sep 17 00:00:00 2001 From: Jason Victor Date: Tue, 25 Feb 2025 17:48:14 -0500 Subject: [PATCH 3/4] docstring --- src/shared/transport.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/shared/transport.ts b/src/shared/transport.ts index bcadcf4d..b80e2a51 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -42,5 +42,8 @@ export interface Transport { */ onmessage?: (message: JSONRPCMessage) => void; + /** + * The session ID generated for this connection. + */ sessionId?: string; } From c396b4af1daea1afd144514348d75d76c93a4d86 Mon Sep 17 00:00:00 2001 From: Jason Victor Date: Tue, 25 Feb 2025 17:57:45 -0500 Subject: [PATCH 4/4] added test for session id using in memory transport --- src/inMemory.ts | 1 + src/server/mcp.test.ts | 53 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/src/inMemory.ts b/src/inMemory.ts index 2763f38c..106a9e7e 100644 --- a/src/inMemory.ts +++ b/src/inMemory.ts @@ -11,6 +11,7 @@ export class InMemoryTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage) => void; + sessionId?: string; /** * Creates a pair of linked in-memory transports that can communicate with each other. One should be passed to a Client and one to a Server. diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts index 8f9bfa77..2e91a568 100644 --- a/src/server/mcp.test.ts +++ b/src/server/mcp.test.ts @@ -323,6 +323,59 @@ describe("tool()", () => { mcpServer.tool("tool2", () => ({ content: [] })); }); + test("should pass sessionId to tool callback via RequestHandlerExtra", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client( + { + name: "test client", + version: "1.0", + }, + { + capabilities: { + tools: {}, + }, + }, + ); + + let receivedSessionId: string | undefined; + mcpServer.tool("test-tool", async (extra) => { + receivedSessionId = extra.sessionId; + return { + content: [ + { + type: "text", + text: "Test response", + }, + ], + }; + }); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + // Set a test sessionId on the server transport + serverTransport.sessionId = "test-session-123"; + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + await client.request( + { + method: "tools/call", + params: { + name: "test-tool", + }, + }, + CallToolResultSchema, + ); + + expect(receivedSessionId).toBe("test-session-123"); + }); + test("should allow client to call server tools", async () => { const mcpServer = new McpServer({ name: "test server",