Fix/mcp session status codes

server/src/mcp/server.ts

@@ -13,7 +13,7 @@
  * Protocol versions: 2025-03-26, 2024-11-05
  */
 
-import { createServer, IncomingMessage, ServerResponse } from "http";
+import { createServer, IncomingMessage, Server, ServerResponse } from "http";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { config } from "./config";
@@ -96,6 +96,22 @@ function sendJson(res: ServerResponse, status: number, data: unknown): void {
   res.end(JSON.stringify(data));
 }
 
+/**
+ * Send 404 for an unknown or expired session.
+ * Per the MCP spec, clients must re-initialize when they receive this.
+ * Matches the SDK's own session-not-found response.
+ */
+function sendSessionNotFound(res: ServerResponse): void {
+  sendJson(res, 404, {
+    jsonrpc: "2.0",
+    error: {
+      code: -32001,
+      message: "Session not found",
+    },
+    id: null,
+  });
+}
+
 /**
  * Parse request body as JSON
  */
@@ -122,11 +138,32 @@ interface SessionContext {
 }
 const sessions = new Map<string, SessionContext>();
 
+// sessionStore owns TTL and per-key limits. Close the transport whenever
+// it drops a session, so the two maps never drift apart.
+sessionStore.onDelete((sessionId) => {
+  const sessionCtx = sessions.get(sessionId);
+  if (!sessionCtx) return;
+
+  sessions.delete(sessionId);
+  sessionCtx.transport.close().catch((error) => {
+    console.error("[MCP] Error closing transport:", error);
+  });
+});
+
 /**
- * Start the MCP server
+ * Look up a live session. Goes through sessionStore so each request
+ * refreshes the TTL and an expired session is dropped.
  */
-export function startMcpServer(): void {
-  const server = createServer(
+function getActiveSession(sessionId: string): SessionContext | undefined {
+  if (!sessionStore.validate(sessionId)) return undefined;
+  return sessions.get(sessionId);
+}
+
+/**
+ * Create the MCP HTTP server without starting it
+ */
+export function createMcpHttpServer(): Server {
+  return createServer(
     async (req: IncomingMessage, res: ServerResponse) => {
       // CORS headers
       res.setHeader("Access-Control-Allow-Origin", "*");
@@ -198,11 +235,7 @@ export function startMcpServer(): void {
         // Handle DELETE - session termination
         if (req.method === "DELETE") {
           if (sessionId) {
-            const sessionCtx = sessions.get(sessionId);
-            if (sessionCtx) {
-              await sessionCtx.transport.close();
-              sessions.delete(sessionId);
-            }
+            // The onDelete listener closes the transport
             sessionStore.delete(sessionId);
           }
           res.writeHead(204);
@@ -216,24 +249,15 @@ export function startMcpServer(): void {
             // Get or create transport for this session
             let transport: StreamableHTTPServerTransport;
 
-            if (sessionId && sessions.has(sessionId)) {
+            if (sessionId) {
               // Existing session - reuse transport
-              const sessionCtx = sessions.get(sessionId);
+              const sessionCtx = getActiveSession(sessionId);
               if (!sessionCtx) {
-                throw new Error("Session not found despite has() check");
+                // Unknown, expired, or evicted session
+                sendSessionNotFound(res);
+                return;
               }
               transport = sessionCtx.transport;
-            } else if (sessionId) {
-              // Session ID provided but transport not found - session expired
-              sendJson(res, 400, {
-                jsonrpc: "2.0",
-                error: {
-                  code: -32000,
-                  message: "Bad Request: Session not found or expired",
-                },
-                id: null,
-              });
-              return;
             } else {
               // No session ID - this should be an initialize request
               // Create new MCP server instance for this session
@@ -283,16 +307,19 @@ export function startMcpServer(): void {
 
         // Handle GET - SSE streaming (for server-initiated messages)
         if (req.method === "GET") {
-          if (!sessionId || !sessions.has(sessionId)) {
-            sendJson(res, 400, {
-              error: "Session ID required for SSE connection",
+          // No session: we offer no standalone SSE stream, so per the MCP
+          // spec answer 405 rather than 400
+          if (!sessionId) {
+            res.setHeader("Allow", "POST, DELETE");
+            sendJson(res, 405, {
+              error: "SSE stream requires an initialized session",
             });
             return;
           }
 
-          const sessionCtx = sessions.get(sessionId);
+          const sessionCtx = getActiveSession(sessionId);
           if (!sessionCtx) {
-            sendJson(res, 400, { error: "Session not found" });
+            sendSessionNotFound(res);
             return;
           }
           await sessionCtx.transport.handleRequest(req, res);
@@ -308,6 +335,13 @@ export function startMcpServer(): void {
       sendJson(res, 404, { error: "Not found" });
     },
   );
+}
+
+/**
+ * Start the MCP server
+ */
+export function startMcpServer(): void {
+  const server = createMcpHttpServer();
 
   server.listen(config.port, config.host, () => {
     console.log(

server/src/mcp/sessions.ts

@@ -14,11 +14,14 @@ interface Session {
   initialized: boolean;
 }
 
+type SessionDeleteListener = (sessionId: string) => void;
+
 /**
  * Simple in-memory session store with TTL and limits
  */
 export class SessionStore {
   private sessions = new Map<string, Session>();
+  private deleteListeners: SessionDeleteListener[] = [];
   private cleanupInterval: ReturnType<typeof setInterval> | null = null;
 
   constructor() {
@@ -92,18 +95,32 @@ export class SessionStore {
     return true;
   }
 
+  /**
+   * Register a callback that runs whenever a session is removed
+   * (explicit delete, TTL expiry, or per-key eviction)
+   */
+  onDelete(listener: SessionDeleteListener): void {
+    this.deleteListeners.push(listener);
+  }
+
   /**
    * Delete a session
    */
   delete(sessionId: string): boolean {
     const deleted = this.sessions.delete(sessionId);
-    if (deleted && config.debug) {
+    if (!deleted) return false;
+
+    if (config.debug) {
       // Sanitize sessionId before logging to prevent log injection (CWE-117)
       // eslint-disable-next-line no-control-regex
       const safeId = sessionId.replace(/[\r\n\t\x00-\x1f\x7f]/g, "_");
       console.log(`[MCP] Session deleted: ${safeId}`);
     }
-    return deleted;
+
+    for (const listener of this.deleteListeners) {
+      listener(sessionId);
+    }
+    return true;
   }
 
   /**
@@ -131,7 +148,7 @@ export class SessionStore {
 
     for (const [id, session] of this.sessions) {
       if (now - session.lastAccessedAt > config.sessionTtlMs) {
-        this.sessions.delete(id);
+        this.delete(id);
         cleaned++;
       }
     }