diff --git a/src/handlers/shttp.ts b/src/handlers/shttp.ts index c725882..2a7ce52 100644 --- a/src/handlers/shttp.ts +++ b/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) | 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(); - } - }); } }