Refactor shttp handler to support user-based session ownership

- Add userId extraction from auth info
- Implement session ownership validation
- Remove DELETE endpoint handling (moved elsewhere)
- Simplify transport lifecycle management
- Add proper 401 responses for unauthorized access

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

src/handlers/shttp.ts

@@ -1,7 +1,7 @@
 import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { Request, Response } from "express";
-import { getFirstShttpTransport, getShttpTransport, isLive, shutdownSession, startServerListeningToRedis } from "../services/redisTransport.js";
+import { getShttpTransport, isSessionOwnedBy, redisRelayToMcpServer, ServerRedisTransport, setSessionOwner, shutdownSession } from "../services/redisTransport.js";
 import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import { randomUUID } from "crypto";
 import { createMcpServer } from "../services/mcp.js";
@@ -16,101 +16,74 @@ declare module "express-serve-static-core" {
   }
 }
 
+function getUserIdFromAuth(auth?: AuthInfo): string | null {
+  return auth?.extra?.userId as string || null;
+}
+
 export async function handleStreamableHTTP(req: Request, res: Response) {
   let shttpTransport: StreamableHTTPServerTransport | undefined = undefined;
-  let cleanup: (() => Promise<void>) | undefined = undefined;
+
+  res.on('finish', async () => {
+    await shttpTransport?.close();
+  });
+
   try {
-    // Handle DELETE requests for session termination
-    if (req.method === 'DELETE') {
-      const sessionId = req.headers['mcp-session-id'] as string | undefined;
-      
-      if (!sessionId) {
-        res.status(400).json({
-          jsonrpc: '2.0',
-          error: {
-            code: -32000,
-            message: 'Bad Request: Session ID required for DELETE request',
-          },
-          id: null,
-        });
-        return;
-      }
-
-      // Check if session exists
-      if (!(await isLive(sessionId))) {
-        res.status(404).json({
-          jsonrpc: '2.0',
-          error: {
-            code: -32001,
-            message: 'Session not found',
-          },
-          id: null,
-        });
-        return;
-      }
-
-      // Send shutdown control message to terminate the session
-      await shutdownSession(sessionId);
-
-      // Respond with success
-      res.status(200).json({
-        jsonrpc: '2.0',
-        result: { status: 'Session terminated successfully' },
-        id: null,
-      });
-      return;
-    }
-
     // Check for existing session ID
     const sessionId = req.headers['mcp-session-id'] as string | undefined;
+    const userId = getUserIdFromAuth(req.auth);
 
-    if (sessionId && await isLive(sessionId)) {
-      // Reuse existing transport
-      ({ shttpTransport, cleanup } = await getShttpTransport(sessionId));
-    } else if (!sessionId && isInitializeRequest(req.body)) {
-      // New initialization request - use JSON response mode
-      const sessionId = randomUUID();
-      
-      const server = createMcpServer();
-      
-      await startServerListeningToRedis(server, sessionId);
-      
-      ({ shttpTransport, cleanup } = await getFirstShttpTransport(sessionId));
-    } else {
-      // Invalid request - no session ID or not initialization request
-      res.status(400).json({
-        jsonrpc: '2.0',
-        error: {
-          code: -32000,
-          message: 'Bad Request: No valid session ID provided',
-        },
-        id: null,
-      });
+    // if no userid, return 401, we shouldn't get here ideally
+    if (!userId) {
+      res.status(401)
       return;
     }
 
+    // incorrect session for the authed user, return 401
+    if (sessionId) {
+      if (!(await isSessionOwnedBy(sessionId, userId))) {
+        res.status(401)
+        return;
+      }
+      // Reuse existing transport for owned session
+      shttpTransport = await getShttpTransport(sessionId);
+    } else if (isInitializeRequest(req.body)) {
+      // New initialization request - use JSON response mode
+      const onsessioninitialized = async (sessionId: string) => {
+        const { server, cleanup: mcpCleanup } = createMcpServer();
+
+        const serverRedisTransport = new ServerRedisTransport(sessionId);
+        serverRedisTransport.onclose = mcpCleanup;
+        await server.connect(serverRedisTransport)
+      
+        // Set session ownership
+        await setSessionOwner(sessionId, userId);
+      }
+
+      const onsessionclosed = async (sessionId: string) => {
+        console.log(`Session ${sessionId} closing`);
+        void shutdownSession(sessionId);
+        console.log(`Session ${sessionId} closed`);
+      }
+
+      const sessionId = randomUUID();
+      shttpTransport = new StreamableHTTPServerTransport({
+        sessionIdGenerator: () => sessionId,
+        onsessionclosed, // dev build only
+        onsessioninitialized,
+      });
+      shttpTransport.onclose = redisRelayToMcpServer(sessionId, shttpTransport);
+    } else {
+      // Invalid request - no session ID and not initialization request
+      res.status(400)
+      return;
+    }
     // Handle the request with existing transport - no need to reconnect
     await shttpTransport.handleRequest(req, res, req.body);
   } catch (error) {
     console.error('Error handling MCP request:', error);
     
     if (!res.headersSent) {
-      res.status(500).json({
-        jsonrpc: '2.0',
-        error: {
-          code: -32603,
-          message: 'Internal server error',
-        },
-        id: null,
-      });
+      res.status(500)
     }
-  } finally {
-    // Set up cleanup when response is complete
-    res.on('finish', async () => {
-      await shttpTransport?.close();
-      if (cleanup) {
-        await cleanup();
-      }
-    });
   }
 }