feat: add session management for streamableHttp [MCP-52]

fmenezes · fmenezes · commit 2be0d17c9421 · 2025-07-16T18:34:49.000+01:00
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -19,6 +19,7 @@
       "name": "Launch Program",
       "skipFiles": ["<node_internals>/**"],
       "program": "${workspaceFolder}/dist/index.js",
+      "args": ["--transport", "http", "--loggers", "stderr", "mcp"],
       "preLaunchTask": "tsc: build - tsconfig.build.json",
       "outFiles": ["${workspaceFolder}/dist/**/*.js"]
     }
diff --git a/src/common/logger.ts b/src/common/logger.ts
@@ -40,12 +40,9 @@ export const LogId = {
     toolUpdateFailure: mongoLogId(1_005_001),
 
     streamableHttpTransportStarted: mongoLogId(1_006_001),
-    streamableHttpTransportStartFailure: mongoLogId(1_006_002),
-    streamableHttpTransportSessionInitialized: mongoLogId(1_006_003),
-    streamableHttpTransportRequestFailure: mongoLogId(1_006_004),
-    streamableHttpTransportCloseRequested: mongoLogId(1_006_005),
-    streamableHttpTransportCloseSuccess: mongoLogId(1_006_006),
-    streamableHttpTransportCloseFailure: mongoLogId(1_006_007),
+    streamableHttpTransportSessionCloseFailure: mongoLogId(1_006_002),
+    streamableHttpTransportRequestFailure: mongoLogId(1_006_003),
+    streamableHttpTransportCloseFailure: mongoLogId(1_006_004),
 } as const;
 
 export abstract class LoggerBase {
diff --git a/src/common/sessionStore.ts b/src/common/sessionStore.ts
@@ -0,0 +1,44 @@
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import logger, { LogId } from "./logger.js";
+
+export class SessionStore {
+    private sessions: Map<string, StreamableHTTPServerTransport> = new Map();
+
+    getSession(sessionId: string): StreamableHTTPServerTransport | undefined {
+        return this.sessions.get(sessionId);
+    }
+
+    setSession(sessionId: string, transport: StreamableHTTPServerTransport): void {
+        if (this.sessions.has(sessionId)) {
+            throw new Error(`Session ${sessionId} already exists`);
+        }
+        this.sessions.set(sessionId, transport);
+    }
+
+    async closeSession(sessionId: string, closeTransport: boolean = true): Promise<void> {
+        if (!this.sessions.has(sessionId)) {
+            throw new Error(`Session ${sessionId} not found`);
+        }
+        if (closeTransport) {
+            const transport = this.sessions.get(sessionId);
+            if (!transport) {
+                throw new Error(`Session ${sessionId} not found`);
+            }
+            try {
+                await transport.close();
+            } catch (error) {
+                logger.error(
+                    LogId.streamableHttpTransportSessionCloseFailure,
+                    "streamableHttpTransport",
+                    `Error closing transport ${sessionId}: ${error instanceof Error ? error.message : String(error)}`
+                );
+            }
+        }
+        this.sessions.delete(sessionId);
+    }
+
+    async closeAllSessions(): Promise<void> {
+        await Promise.all(Array.from(this.sessions.values()).map((transport) => transport.close()));
+        this.sessions.clear();
+    }
+}
diff --git a/src/transports/streamableHttp.ts b/src/transports/streamableHttp.ts
@@ -1,90 +1,132 @@
 import express from "express";
 import http from "http";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import { TransportRunnerBase } from "./base.js";
 import { config } from "../common/config.js";
 import logger, { LogId } from "../common/logger.js";
+import { randomUUID } from "crypto";
+import { SessionStore } from "../common/sessionStore.js";
 
 const JSON_RPC_ERROR_CODE_PROCESSING_REQUEST_FAILED = -32000;
-const JSON_RPC_ERROR_CODE_METHOD_NOT_ALLOWED = -32601;
+const JSON_RPC_ERROR_CODE_SESSION_ID_REQUIRED = -32001;
+const JSON_RPC_ERROR_CODE_SESSION_NOT_FOUND = -32002;
+const JSON_RPC_ERROR_CODE_INVALID_REQUEST = -32003;
 
 function promiseHandler(
     fn: (req: express.Request, res: express.Response, next: express.NextFunction) => Promise<void>
 ) {
     return (req: express.Request, res: express.Response, next: express.NextFunction) => {
-        fn(req, res, next).catch(next);
+        fn(req, res, next).catch((error) => {
+            logger.error(
+                LogId.streamableHttpTransportRequestFailure,
+                "streamableHttpTransport",
+                `Error handling request: ${error instanceof Error ? error.message : String(error)}`
+            );
+            res.status(400).json({
+                jsonrpc: "2.0",
+                error: {
+                    code: JSON_RPC_ERROR_CODE_PROCESSING_REQUEST_FAILED,
+                    message: `failed to handle request`,
+                    data: error instanceof Error ? error.message : String(error),
+                },
+            });
+        });
     };
 }
 
 export class StreamableHttpRunner extends TransportRunnerBase {
     private httpServer: http.Server | undefined;
+    private sessionStore: SessionStore = new SessionStore();
 
     async start() {
         const app = express();
         app.enable("trust proxy"); // needed for reverse proxy support
-        app.use(express.urlencoded({ extended: true }));
         app.use(express.json());
 
+        const handleRequest = async (req: express.Request, res: express.Response) => {
+            const sessionId = req.headers["mcp-session-id"] as string;
+            if (!sessionId) {
+                res.status(400).json({
+                    jsonrpc: "2.0",
+                    error: {
+                        code: JSON_RPC_ERROR_CODE_SESSION_ID_REQUIRED,
+                        message: `session id is required`,
+                    },
+                });
+                return;
+            }
+            const transport = this.sessionStore.getSession(sessionId);
+            if (!transport) {
+                res.status(404).json({
+                    jsonrpc: "2.0",
+                    error: {
+                        code: JSON_RPC_ERROR_CODE_SESSION_NOT_FOUND,
+                        message: `session not found`,
+                    },
+                });
+                return;
+            }
+            await transport.handleRequest(req, res, req.body);
+        };
+
         app.post(
             "/mcp",
             promiseHandler(async (req: express.Request, res: express.Response) => {
-                const transport = new StreamableHTTPServerTransport({
-                    sessionIdGenerator: undefined,
-                });
+                const sessionId = req.headers["mcp-session-id"] as string;
+                if (sessionId) {
+                    await handleRequest(req, res);
+                    return;
+                }
 
-                const server = this.setupServer();
+                if (!isInitializeRequest(req.body)) {
+                    res.status(400).json({
+                        jsonrpc: "2.0",
+                        error: {
+                            code: JSON_RPC_ERROR_CODE_INVALID_REQUEST,
+                            message: `invalid request`,
+                        },
+                    });
+                    return;
+                }
 
-                await server.connect(transport);
+                const server = this.setupServer();
+                const transport = new StreamableHTTPServerTransport({
+                    sessionIdGenerator: () => randomUUID().toString(),
+                    onsessioninitialized: (sessionId) => {
+                        this.sessionStore.setSession(sessionId, transport);
+                    },
+                    onsessionclosed: async (sessionId) => {
+                        try {
+                            await this.sessionStore.closeSession(sessionId, false);
+                        } catch (error) {
+                            logger.error(
+                                LogId.streamableHttpTransportSessionCloseFailure,
+                                "streamableHttpTransport",
+                                `Error closing session: ${error instanceof Error ? error.message : String(error)}`
+                            );
+                        }
+                    },
+                });
 
-                res.on("close", () => {
-                    Promise.all([transport.close(), server.close()]).catch((error: unknown) => {
+                transport.onclose = () => {
+                    server.close().catch((error) => {
                         logger.error(
                             LogId.streamableHttpTransportCloseFailure,
                             "streamableHttpTransport",
                             `Error closing server: ${error instanceof Error ? error.message : String(error)}`
                         );
                     });
-                });
+                };
 
-                try {
-                    await transport.handleRequest(req, res, req.body);
-                } catch (error) {
-                    logger.error(
-                        LogId.streamableHttpTransportRequestFailure,
-                        "streamableHttpTransport",
-                        `Error handling request: ${error instanceof Error ? error.message : String(error)}`
-                    );
-                    res.status(400).json({
-                        jsonrpc: "2.0",
-                        error: {
-                            code: JSON_RPC_ERROR_CODE_PROCESSING_REQUEST_FAILED,
-                            message: `failed to handle request`,
-                            data: error instanceof Error ? error.message : String(error),
-                        },
-                    });
-                }
+                await server.connect(transport);
+
+                await transport.handleRequest(req, res, req.body);
             })
         );
 
-        app.get("/mcp", (req: express.Request, res: express.Response) => {
-            res.status(405).json({
-                jsonrpc: "2.0",
-                error: {
-                    code: JSON_RPC_ERROR_CODE_METHOD_NOT_ALLOWED,
-                    message: `method not allowed`,
-                },
-            });
-        });
-
-        app.delete("/mcp", (req: express.Request, res: express.Response) => {
-            res.status(405).json({
-                jsonrpc: "2.0",
-                error: {
-                    code: JSON_RPC_ERROR_CODE_METHOD_NOT_ALLOWED,
-                    message: `method not allowed`,
-                },
-            });
-        });
+        app.get("/mcp", promiseHandler(handleRequest));
+        app.delete("/mcp", promiseHandler(handleRequest));
 
         this.httpServer = await new Promise<http.Server>((resolve, reject) => {
             const result = app.listen(config.httpPort, config.httpHost, (err?: Error) => {
diff --git a/tests/integration/transports/stdio.test.ts b/tests/integration/transports/stdio.test.ts
@@ -1,70 +1,40 @@
-import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
 import { describe, expect, it, beforeAll, afterAll } from "vitest";
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 
 describe("StdioRunner", () => {
     describe("client connects successfully", () => {
-        let client: StdioClientTransport;
+        let client: Client;
+        let transport: StdioClientTransport;
         beforeAll(async () => {
-            client = new StdioClientTransport({
+            transport = new StdioClientTransport({
                 command: "node",
                 args: ["dist/index.js"],
                 env: {
                     MDB_MCP_TRANSPORT: "stdio",
                 },
             });
-            await client.start();
+            client = new Client({
+                name: "test",
+                version: "0.0.0",
+            });
+            await client.connect(transport);
         });
 
         afterAll(async () => {
             await client.close();
+            await transport.close();
         });
 
         it("handles requests and sends responses", async () => {
-            let fixedResolve: ((value: JSONRPCMessage) => void) | undefined = undefined;
-            const messagePromise = new Promise<JSONRPCMessage>((resolve) => {
-                fixedResolve = resolve;
-            });
-
-            client.onmessage = (message: JSONRPCMessage) => {
-                fixedResolve?.(message);
-            };
-
-            await client.send({
-                jsonrpc: "2.0",
-                id: 1,
-                method: "tools/list",
-                params: {
-                    _meta: {
-                        progressToken: 1,
-                    },
-                },
-            });
-
-            const message = (await messagePromise) as {
-                jsonrpc: string;
-                id: number;
-                result: {
-                    tools: {
-                        name: string;
-                        description: string;
-                    }[];
-                };
-                error?: {
-                    code: number;
-                    message: string;
-                };
-            };
+            const response = await client.listTools();
+            expect(response).toBeDefined();
+            expect(response.tools).toBeDefined();
+            expect(response.tools).toHaveLength(20);
 
-            expect(message.jsonrpc).toBe("2.0");
-            expect(message.id).toBe(1);
-            expect(message.result).toBeDefined();
-            expect(message.result?.tools).toBeDefined();
-            expect(message.result?.tools.length).toBeGreaterThan(0);
-            const tools = message.result?.tools;
-            tools.sort((a, b) => a.name.localeCompare(b.name));
-            expect(tools[0]?.name).toBe("aggregate");
-            expect(tools[0]?.description).toBe("Run an aggregation against a MongoDB collection");
+            const sortedTools = response.tools.sort((a, b) => a.name.localeCompare(b.name));
+            expect(sortedTools[0]?.name).toBe("aggregate");
+            expect(sortedTools[0]?.description).toBe("Run an aggregation against a MongoDB collection");
         });
     });
 });
diff --git a/tests/integration/transports/streamableHttp.test.ts b/tests/integration/transports/streamableHttp.test.ts

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@`
`19`	`19`	`"name": "Launch Program",`
`20`	`20`	`"skipFiles": ["<node_internals>/**"],`
`21`	`21`	`"program": "${workspaceFolder}/dist/index.js",`
	`22`	`+ "args": ["--transport", "http", "--loggers", "stderr", "mcp"],`
`22`	`23`	`"preLaunchTask": "tsc: build - tsconfig.build.json",`
`23`	`24`	`"outFiles": ["${workspaceFolder}/dist/*/.js"]`
`24`	`25`	`}`