From 8579d446aaeb755f3bace358ee8d6dce4c837c6a Mon Sep 17 00:00:00 2001 From: Jerome Date: Mon, 30 Jun 2025 20:29:41 +0100 Subject: [PATCH] Add comprehensive logging to debug Redis transport issues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add detailed logging to Redis transport lifecycle - Log session validation and message flow - Track MCP server request handling - Add diagnostics for streamable HTTP transport connections 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- .vscode/launch.json | 18 ++ src/handlers/common.ts | 57 +++++ src/handlers/mcp.ts | 344 ---------------------------- src/handlers/shttp.ts | 115 ++++++++++ src/handlers/sse.ts | 102 +++++++++ src/index.ts | 17 +- src/redis.ts | 20 ++ src/services/auth.ts | 4 +- src/services/mcp.ts | 12 + src/services/messageDelivery.ts | 102 --------- src/services/redisTransport.ts | 247 ++++++++++++++++++++ src/services/sessionManager.test.ts | 160 ------------- src/services/sessionManager.ts | 157 ------------- 13 files changed, 584 insertions(+), 771 deletions(-) create mode 100644 .vscode/launch.json create mode 100644 src/handlers/common.ts delete mode 100644 src/handlers/mcp.ts create mode 100644 src/handlers/shttp.ts create mode 100644 src/handlers/sse.ts delete mode 100644 src/services/messageDelivery.ts create mode 100644 src/services/redisTransport.ts delete mode 100644 src/services/sessionManager.test.ts delete mode 100644 src/services/sessionManager.ts diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..435c9f3 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,18 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "Debug MCP Server (tsx watch)", + "type": "node", + "request": "launch", + "program": "${workspaceFolder}/node_modules/.bin/tsx", + "args": ["watch", "src/index.ts"], + "console": "integratedTerminal", + "restart": true, + "skipFiles": ["/**"], + "env": { + "NODE_ENV": "development" + } + } + ] +} \ No newline at end of file diff --git a/src/handlers/common.ts b/src/handlers/common.ts new file mode 100644 index 0000000..02b52b0 --- /dev/null +++ b/src/handlers/common.ts @@ -0,0 +1,57 @@ +import { NextFunction, Request, Response } from "express"; +import { withContext } from "../context.js"; +import { readMcpInstallation } from "../services/auth.js"; + +import { JSONRPCError, JSONRPCNotification, JSONRPCRequest, JSONRPCResponse } from "@modelcontextprotocol/sdk/types.js"; + +export function logMcpMessage( + message: JSONRPCError | JSONRPCNotification | JSONRPCRequest | JSONRPCResponse, + sessionId: string, +) { + // check if message has a method field + if ("method" in message) { + if (message.method === "tools/call") { + console.info( + `[session ${sessionId}] Processing ${message.method}, for tool ${message.params?.name}`, + ); + } else { + console.info( + `[session ${sessionId}] Processing ${message.method} method`, + ); + } + } else if ("error" in message) { + console.warn( + `[session ${sessionId}] Received error message: ${message.error.message}, ${message.error.code}`, + ) + } +} + + +export async function authContext( + req: Request, + res: Response, + next: NextFunction, +) { + const authInfo = req.auth + + if (!authInfo) { + res.set("WWW-Authenticate", 'Bearer error="invalid_token"'); + res.status(401).json({ error: "Invalid access token" }); + return; + } + + const token = authInfo.token; + + // Load UpstreamInstallation based on the access token + const mcpInstallation = await readMcpInstallation(token); + if (!mcpInstallation) { + res.set("WWW-Authenticate", 'Bearer error="invalid_token"'); + res.status(401).json({ error: "Invalid access token" }); + return; + } + + // Wrap the rest of the request handling in the context + withContext({ mcpAccessToken: token, fakeUpstreamInstallation: mcpInstallation.fakeUpstreamInstallation }, () => + next(), + ); +} diff --git a/src/handlers/mcp.ts b/src/handlers/mcp.ts deleted file mode 100644 index 8a4b531..0000000 --- a/src/handlers/mcp.ts +++ /dev/null @@ -1,344 +0,0 @@ -import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js"; -import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; -import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js"; -import contentType from "content-type"; -import { NextFunction, Request, Response } from "express"; -import getRawBody from "raw-body"; -import { randomUUID } from "node:crypto"; -import { readMcpInstallation } from "../services/auth.js"; -import { withContext } from "../context.js"; -import { createMcpServer } from "../services/mcp.js"; -import { redisClient } from "../redis.js"; -import { SessionManager } from "../services/sessionManager.js"; -import { MessageDelivery } from "../services/messageDelivery.js"; -import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js"; - -const MAXIMUM_MESSAGE_SIZE = "4mb"; - -declare module "express-serve-static-core" { - interface Request { - /** - * Information about the validated access token, if the `requireBearerAuth` middleware was used. - */ - auth?: AuthInfo; - } -} - -export async function authContext( - req: Request, - res: Response, - next: NextFunction, -) { - const authInfo = req.auth - - if (!authInfo) { - res.set("WWW-Authenticate", 'Bearer error="invalid_token"'); - res.status(401).json({ error: "Invalid access token" }); - return; - } - - const token = authInfo.token; - - // Load UpstreamInstallation based on the access token - const mcpInstallation = await readMcpInstallation(token); - if (!mcpInstallation) { - res.set("WWW-Authenticate", 'Bearer error="invalid_token"'); - res.status(401).json({ error: "Invalid access token" }); - return; - } - - // Wrap the rest of the request handling in the context - withContext({ mcpAccessToken: token, fakeUpstreamInstallation: mcpInstallation.fakeUpstreamInstallation }, () => - next(), - ); -} - -function redisChannelForSession(sessionId: string): string { - return `mcp:${sessionId}`; -} - -export async function handleSSEConnection(req: Request, res: Response) { - const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer(); - const transport = new SSEServerTransport("/message", res); - console.info(`[session ${transport.sessionId}] Received MCP SSE connection`); - - const redisCleanup = await redisClient.createSubscription( - redisChannelForSession(transport.sessionId), - (json) => { - const message = JSON.parse(json); - - if (message.method) { - if (message.method === "tools/call") { - console.info( - `[session ${transport.sessionId}] Processing ${message.method}, for tool ${message.params?.name}`, - ); - } else { - console.info( - `[session ${transport.sessionId}] Processing ${message.method} method`, - ); - } - } else if (message.error) { - console.warn( - `[session ${transport.sessionId}] Received error message: ${message.error.message}, ${message.error.code}`, - ) - } - transport.handleMessage(message).catch((error) => { - console.error( - `[session ${transport.sessionId}] Error handling message:`, - error, - ); - }); - }, - (error) => { - console.error( - `[session ${transport.sessionId}] Disconnecting due to error in Redis subscriber:`, - error, - ); - transport - .close() - .catch((error) => - console.error( - `[session ${transport.sessionId}] Error closing transport:`, - error, - ), - ); - }, - ); - - const cleanup = () => { - void mcpCleanup(); - redisCleanup().catch((error) => - console.error( - `[session ${transport.sessionId}] Error disconnecting Redis subscriber:`, - error, - ), - ); - } - - // Clean up Redis subscription when the connection closes - mcpServer.onclose = cleanup - - console.info(`[session ${transport.sessionId}] Listening on Redis channel`); - await mcpServer.connect(transport); -} - -export async function handleMessage(req: Request, res: Response) { - const sessionId = req.query.sessionId; - let body: string; - try { - if (typeof sessionId !== "string") { - throw new Error("Only one sessionId allowed"); - } - - const ct = contentType.parse(req.headers["content-type"] ?? ""); - if (ct.type !== "application/json") { - throw new Error(`Unsupported content-type: ${ct}`); - } - - body = await getRawBody(req, { - limit: MAXIMUM_MESSAGE_SIZE, - encoding: ct.parameters.charset ?? "utf-8", - }); - } catch (error) { - res.status(400).json(error); - console.error("Bad POST request:", error); - return; - } - await redisClient.publish(redisChannelForSession(sessionId), body); - res.status(202).end(); -} - -// Initialize session and message delivery services -const sessionManager = new SessionManager(redisClient); -const messageDelivery = new MessageDelivery(redisClient, sessionManager); - -/** - * Handler for the new streamable HTTP transport (/mcp endpoint) - * Supports GET, POST, and DELETE methods with Redis-based session management - */ -export async function handleStreamableHTTP(req: Request, res: Response) { - const method = req.method; - const sessionId = req.headers['mcp-session-id'] as string | undefined; - - try { - if (method === 'POST') { - await handleStreamableHTTPPost(req, res, sessionId); - } else if (method === 'GET') { - await handleStreamableHTTPGet(req, res, sessionId); - } else if (method === 'DELETE') { - await handleStreamableHTTPDelete(req, res, sessionId); - } else { - res.writeHead(405, { "Allow": "GET, POST, DELETE" }).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32000, message: "Method not allowed." }, - id: null - })); - } - } catch (error) { - console.error(`Error handling streamable HTTP ${method} request:`, error); - if (!res.headersSent) { - res.writeHead(500).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32603, message: "Internal server error" }, - id: null - })); - } - } -} - -/** - * Handle POST requests (initialization or message sending) - */ -async function handleStreamableHTTPPost(req: Request, res: Response, sessionId: string | undefined) { - // Parse request body - const ct = req.headers["content-type"]; - if (!ct || !ct.includes("application/json")) { - res.writeHead(415).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32000, message: "Unsupported Media Type: Content-Type must be application/json" }, - id: null - })); - return; - } - - const parsedCt = contentType.parse(ct); - const body = await getRawBody(req, { - limit: MAXIMUM_MESSAGE_SIZE, - encoding: parsedCt.parameters.charset ?? "utf-8", - }); - const parsedBody = JSON.parse(body.toString()); - - if (isInitializeRequest(parsedBody) && !sessionId) { - // New session initialization - const newSessionId = randomUUID(); - const authInfo = req.auth; - - console.log(`[streamable-http] Initializing new session: ${newSessionId}`); - - // Create session in Redis - await sessionManager.createSession(newSessionId, authInfo?.clientId || 'unknown'); - - // Create transport with the new session ID - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => newSessionId, - }); - - // Create MCP server and connect transport - const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer(); - await mcpServer.connect(transport); - - // Handle the initialization request - await transport.handleRequest(req, res, parsedBody); - - // Clean up after request completes - res.on('close', mcpCleanup); - - } else if (sessionId) { - // Existing session - validate and handle request - const sessionValid = await sessionManager.refreshSession(sessionId); - if (!sessionValid) { - res.writeHead(404).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32001, message: "Session not found" }, - id: null - })); - return; - } - - console.log(`[streamable-http] Handling request for existing session: ${sessionId}`); - - // Create ephemeral transport for this request - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => sessionId, - }); - - // Create MCP server and connect transport - const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer(); - await mcpServer.connect(transport); - - // Handle the request - await transport.handleRequest(req, res, parsedBody); - - // Clean up after request completes - res.on('close', mcpCleanup); - - } else { - // Invalid request - no session ID and not initialization - res.writeHead(400).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32000, message: "Bad Request: Session ID required for non-initialization requests" }, - id: null - })); - } -} - -/** - * Handle GET requests (SSE stream establishment) - */ -async function handleStreamableHTTPGet(req: Request, res: Response, sessionId: string | undefined) { - if (!sessionId) { - res.writeHead(400).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32000, message: "Bad Request: Session ID required" }, - id: null - })); - return; - } - - // Validate session exists and refresh TTL - const sessionValid = await sessionManager.refreshSession(sessionId); - if (!sessionValid) { - res.writeHead(404).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32001, message: "Session not found" }, - id: null - })); - return; - } - - console.log(`[streamable-http] Establishing SSE stream for session: ${sessionId}`); - - // Create transport for SSE stream - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => sessionId, - }); - - // Mark connection as active - await sessionManager.setConnectionState(sessionId, true); - - // Deliver any buffered messages first - await messageDelivery.deliverBufferedMessages(sessionId, transport); - - // Set up Redis subscription for live messages - const redisCleanup = await messageDelivery.setupRedisSubscription(sessionId, transport); - - // Handle connection cleanup when client disconnects - res.on('close', async () => { - console.log(`[streamable-http] SSE connection closed for session: ${sessionId}`); - await sessionManager.setConnectionState(sessionId, false); - await redisCleanup(); - }); - - // Start the SSE stream - await transport.handleRequest(req, res); -} - -/** - * Handle DELETE requests (session termination) - */ -async function handleStreamableHTTPDelete(req: Request, res: Response, sessionId: string | undefined) { - if (!sessionId) { - res.writeHead(400).end(JSON.stringify({ - jsonrpc: "2.0", - error: { code: -32000, message: "Bad Request: Session ID required" }, - id: null - })); - return; - } - - console.log(`[streamable-http] Deleting session: ${sessionId}`); - - // Delete session from Redis - await sessionManager.deleteSession(sessionId); - res.writeHead(200).end(); -} diff --git a/src/handlers/shttp.ts b/src/handlers/shttp.ts new file mode 100644 index 0000000..79eba9b --- /dev/null +++ b/src/handlers/shttp.ts @@ -0,0 +1,115 @@ +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, startServerListeningToRedis } from "../services/redisTransport.js"; +import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js"; +import { randomUUID } from "crypto"; +import { createMcpServer } from "../services/mcp.js"; +import getRawBody from "raw-body"; + + +declare module "express-serve-static-core" { + interface Request { + /** + * Information about the validated access token, if the `requireBearerAuth` middleware was used. + */ + auth?: AuthInfo; + } +} + +export async function handleStreamableHTTP(req: Request, res: Response) { + console.log('=== handleStreamableHTTP START ==='); + console.log('Received MCP request:', JSON.stringify(req.body, null, 2)); + console.log('Request headers:', JSON.stringify(req.headers, null, 2)); + + let transport: StreamableHTTPServerTransport | undefined = undefined; + try { + // Check for existing session ID + const sessionId = req.headers['mcp-session-id'] as string | undefined; + console.log('Session ID from headers:', sessionId); + + if (sessionId && await isLive(sessionId)) { + console.log('Session is live, reusing existing transport for session:', sessionId); + // Reuse existing transport + transport = await getShttpTransport(sessionId) + console.log('Retrieved transport from Redis for session:', sessionId); + } else if (!sessionId && isInitializeRequest(req.body)) { + console.log('New initialization request detected, creating new session'); + // New initialization request - use JSON response mode + const sessionId = randomUUID(); + console.log('Generated new session ID:', sessionId); + + const server = createMcpServer(); + console.log('Created MCP server instance'); + + startServerListeningToRedis(server.server, sessionId) + console.log('Started server listening to Redis for session:', sessionId); + + transport = await getFirstShttpTransport(sessionId); + console.log('Retrieved first transport for session:', sessionId); + console.log('Transport object:', transport.constructor.name, 'sessionId:', transport.sessionId); + + // Connect the transport to the MCP server BEFORE handling the request + console.log('Connecting transport to MCP server...'); + await server.server.connect(transport); + console.log('Transport connected successfully'); + + console.log('Handling initialization request...'); + await transport.handleRequest(req, res, req.body); + console.log('Initialization request handled successfully'); + console.log('=== handleStreamableHTTP END (initialization) ==='); + return; // Already handled + } else { + console.log('Invalid request - no session ID and not an initialization request'); + console.log('Session ID present:', !!sessionId); + console.log('Is initialize request:', isInitializeRequest(req.body)); + // 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, + }); + console.log('Sent 400 Bad Request response'); + console.log('=== handleStreamableHTTP END (bad request) ==='); + return; + } + + // Handle the request with existing transport - no need to reconnect + console.log('Handling request with existing transport for session:', sessionId); + console.log('Transport object:', transport.constructor.name, 'sessionId:', transport.sessionId); + console.log('Request body:', JSON.stringify(req.body, null, 2)); + await transport.handleRequest(req, res, req.body); + console.log('Request handled successfully'); + } catch (error) { + console.error('=== ERROR in handleStreamableHTTP ==='); + console.error('Error type:', error instanceof Error ? error.constructor.name : typeof error); + console.error('Error message:', error instanceof Error ? error.message : String(error)); + console.error('Error stack:', error instanceof Error ? error.stack : 'No stack trace'); + console.error('Error handling MCP request:', error); + + if (!res.headersSent) { + console.log('Sending 500 Internal Server Error response'); + res.status(500).json({ + jsonrpc: '2.0', + error: { + code: -32603, + message: 'Internal server error', + }, + id: null, + }); + } else { + console.log('Response headers already sent, cannot send error response'); + } + } finally { + // if (transport) { + // console.log('Closing transport in finally block'); + // // Close transports because they are ephemeral in this setup. + // transport.close(); + // console.log('Transport closed'); + // } + console.log('=== handleStreamableHTTP END ==='); + } +} diff --git a/src/handlers/sse.ts b/src/handlers/sse.ts new file mode 100644 index 0000000..8896603 --- /dev/null +++ b/src/handlers/sse.ts @@ -0,0 +1,102 @@ +import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js"; +import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js"; +import contentType from "content-type"; +import { Request, Response } from "express"; +import getRawBody from "raw-body"; +import { redisClient } from "../redis.js"; +import { createMcpServer } from "../services/mcp.js"; +import { logMcpMessage } from "./common.js"; + +const MAXIMUM_MESSAGE_SIZE = "4mb"; + +declare module "express-serve-static-core" { + interface Request { + /** + * Information about the validated access token, if the `requireBearerAuth` middleware was used. + */ + auth?: AuthInfo; + } +} + +function redisChannelForSession(sessionId: string): string { + return `mcp:${sessionId}`; +} + +export async function handleSSEConnection(req: Request, res: Response) { + const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer(); + const transport = new SSEServerTransport("/message", res); + console.info(`[session ${transport.sessionId}] Received MCP SSE connection`); + + const redisCleanup = await redisClient.createSubscription( + redisChannelForSession(transport.sessionId), + (json) => { + // TODO handle DELETE messages + // TODO set timeout to kill the session + + const message = JSON.parse(json); + logMcpMessage(message, transport.sessionId); + transport.handleMessage(message).catch((error) => { + console.error( + `[session ${transport.sessionId}] Error handling message:`, + error, + ); + }); + }, + (error) => { + console.error( + `[session ${transport.sessionId}] Disconnecting due to error in Redis subscriber:`, + error, + ); + transport + .close() + .catch((error) => + console.error( + `[session ${transport.sessionId}] Error closing transport:`, + error, + ), + ); + }, + ); + + const cleanup = () => { + void mcpCleanup(); + redisCleanup().catch((error) => + console.error( + `[session ${transport.sessionId}] Error disconnecting Redis subscriber:`, + error, + ), + ); + } + + // Clean up Redis subscription when the connection closes + mcpServer.onclose = cleanup + + console.info(`[session ${transport.sessionId}] Listening on Redis channel`); + await mcpServer.connect(transport); +} + +export async function handleMessage(req: Request, res: Response) { + const sessionId = req.query.sessionId; + let body: string; + try { + if (typeof sessionId !== "string") { + throw new Error("Only one sessionId allowed"); + } + + const ct = contentType.parse(req.headers["content-type"] ?? ""); + if (ct.type !== "application/json") { + throw new Error(`Unsupported content-type: ${ct}`); + } + + body = await getRawBody(req, { + limit: MAXIMUM_MESSAGE_SIZE, + encoding: ct.parameters.charset ?? "utf-8", + }); + } catch (error) { + res.status(400).json(error); + console.error("Bad POST request:", error); + return; + } + await redisClient.publish(redisChannelForSession(sessionId), body); + res.status(202).end(); +} diff --git a/src/index.ts b/src/index.ts index e2e6dd0..9d67291 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,12 +1,14 @@ +import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js"; +import { AuthRouterOptions, mcpAuthRouter } from "@modelcontextprotocol/sdk/server/auth/router.js"; import cors from "cors"; import express from "express"; -import { BASE_URI, PORT } from "./config.js"; -import { AuthRouterOptions, mcpAuthRouter } from "@modelcontextprotocol/sdk/server/auth/router.js"; import { EverythingAuthProvider } from "./auth/provider.js"; -import { handleMessage, handleSSEConnection, handleStreamableHTTP, authContext } from "./handlers/mcp.js"; -import { handleFakeAuthorizeRedirect, handleFakeAuthorize } from "./handlers/fakeauth.js"; +import { BASE_URI, PORT } from "./config.js"; +import { authContext } from "./handlers/common.js"; +import { handleFakeAuthorize, handleFakeAuthorizeRedirect } from "./handlers/fakeauth.js"; +import { handleStreamableHTTP } from "./handlers/shttp.js"; +import { handleMessage, handleSSEConnection } from "./handlers/sse.js"; import { redisClient } from "./redis.js"; -import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js"; const app = express(); @@ -61,10 +63,13 @@ const sseHeaders = (req: express.Request, res: express.Response, next: express.N const corsOptions = { origin: true, // Allow any origin methods: ['GET', 'POST'], - allowedHeaders: ['Content-Type', 'Authorization', "MCP-Protocol-Version"], + allowedHeaders: ['Content-Type', 'Authorization', "Mcp-Protocol-Version", "Mcp-Protocol-Id"], + exposedHeaders: ["Mcp-Protocol-Version", "Mcp-Protocol-Id"], credentials: true }; + +app.use(express.json()); app.use(logger); // Apply base security headers to all routes diff --git a/src/redis.ts b/src/redis.ts index dda7a27..5ad7d87 100644 --- a/src/redis.ts +++ b/src/redis.ts @@ -15,6 +15,8 @@ export interface RedisClient { connect(): Promise; on(event: string, callback: (error: Error) => void): void; options?: { url: string }; + exists(key: string): Promise; + numsub(key: string): Promise; /** * Creates a pub/sub subscription. Returns a cleanup function to unsubscribe. @@ -48,6 +50,11 @@ export class RedisClientImpl implements RedisClient { ); } + async numsub(key: string): Promise { + const subs = await this.redis.pubSubNumSub(key); + return subs[key] || 0; + } + async get(key: string): Promise { return await this.redis.get(key); } @@ -110,6 +117,11 @@ export class RedisClientImpl implements RedisClient { async publish(channel: string, message: string): Promise { await this.redis.publish(channel, message); } + + async exists(key: string): Promise { + const result = await this.redis.exists(key); + return result > 0; + } } // Export a mutable reference that can be swapped in tests @@ -237,6 +249,14 @@ export class MockRedisClient implements RedisClient { } } + async exists(key: string): Promise { + return this.store.has(key) || this.lists.has(key); + } + + async numsub(key: string): Promise { + return (this.subscribers.get(key) || []).length; + } + clear() { this.store.clear(); this.lists.clear(); diff --git a/src/services/auth.ts b/src/services/auth.ts index c4fa5ee..f9197c0 100644 --- a/src/services/auth.ts +++ b/src/services/auth.ts @@ -243,7 +243,7 @@ export async function exchangeToken( const tokenExchange: TokenExchange = JSON.parse(decoded); if (tokenExchange.alreadyUsed) { - console.error("Duplicate use of authorization code detected; revoking tokens"); + console.error("1Duplicate use of authorization code detected; revoking tokens"); await revokeMcpInstallation(tokenExchange.mcpAccessToken); throw new Error("Duplicate use of authorization code detected; tokens revoked"); } @@ -257,7 +257,7 @@ export async function exchangeToken( if (rereadData !== data) { // Data concurrently changed while we were updating it. This necessarily means a duplicate use. - console.error("Duplicate use of authorization code detected; revoking tokens"); + console.error("2Duplicate use of authorization code detected; revoking tokens"); await revokeMcpInstallation(tokenExchange.mcpAccessToken); throw new Error("Duplicate use of authorization code detected; tokens revoked"); } diff --git a/src/services/mcp.ts b/src/services/mcp.ts index 16f20fc..8d8711e 100644 --- a/src/services/mcp.ts +++ b/src/services/mcp.ts @@ -109,6 +109,18 @@ export const createMcpServer = () => { } ); + // Add logging to track all incoming messages + const originalSetRequestHandler = server.setRequestHandler.bind(server); + server.setRequestHandler = function(schema: any, handler: any) { + const wrappedHandler = async (request: any) => { + console.log(`[MCP Server] Handling request: ${request.method}`); + const result = await handler(request); + console.log(`[MCP Server] Request ${request.method} completed`); + return result; + }; + return originalSetRequestHandler(schema, wrappedHandler); + }; + const subscriptions: Set = new Set(); // Set up update interval for subscribed resources diff --git a/src/services/messageDelivery.ts b/src/services/messageDelivery.ts deleted file mode 100644 index 1fe6cba..0000000 --- a/src/services/messageDelivery.ts +++ /dev/null @@ -1,102 +0,0 @@ -import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js"; -import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; -import { RedisClient } from "../redis.js"; -import { SessionManager } from "./sessionManager.js"; - -/** - * Handles message delivery with buffering support for disconnected clients - */ -export class MessageDelivery { - constructor( - private redisClient: RedisClient, - private sessionManager: SessionManager - ) {} - - /** - * Deliver a message to a session, either directly or via buffering - */ - async deliverMessage(sessionId: string, message: JSONRPCMessage): Promise { - const isConnected = await this.sessionManager.isConnected(sessionId); - - if (isConnected) { - // Direct delivery via existing Redis pub/sub - const redisChannel = this.getRedisChannel(sessionId); - await this.redisClient.publish(redisChannel, JSON.stringify(message)); - console.log(`Message delivered directly to session ${sessionId}`); - } else { - // Buffer the message for later delivery - await this.sessionManager.bufferMessage(sessionId, JSON.stringify(message)); - console.log(`Message buffered for disconnected session ${sessionId}`); - } - } - - /** - * Deliver all buffered messages to a transport and clear the buffer - */ - async deliverBufferedMessages(sessionId: string, transport: StreamableHTTPServerTransport): Promise { - // Get all buffered messages - const bufferedMessages = await this.sessionManager.getBufferedMessages(sessionId); - - if (bufferedMessages.length === 0) { - console.log(`No buffered messages for session ${sessionId}`); - return; - } - - console.log(`Delivering ${bufferedMessages.length} buffered messages to session ${sessionId}`); - - // Deliver buffered messages in order (reverse because lpush adds to front) - for (let i = bufferedMessages.length - 1; i >= 0; i--) { - const message = JSON.parse(bufferedMessages[i]); - try { - await transport.send(message); - } catch (error) { - console.error(`Failed to deliver buffered message to session ${sessionId}:`, error); - // Don't rethrow - continue with other messages - } - } - - // Clear the buffer after successful delivery - await this.sessionManager.clearBufferedMessages(sessionId); - console.log(`Cleared buffered messages for session ${sessionId}`); - } - - /** - * Set up Redis subscription for a session's SSE connection - */ - async setupRedisSubscription( - sessionId: string, - transport: StreamableHTTPServerTransport - ): Promise<() => Promise> { - const redisChannel = this.getRedisChannel(sessionId); - - console.log(`Setting up Redis subscription for session ${sessionId} on channel ${redisChannel}`); - - const redisCleanup = await this.redisClient.createSubscription( - redisChannel, - async (message) => { - const jsonMessage = JSON.parse(message); - try { - await transport.send(jsonMessage); - console.log(`Message sent via SSE for session ${sessionId}`); - } catch (error) { - console.error(`Failed to send message via SSE for session ${sessionId}:`, error); - // Mark connection as disconnected so future messages get buffered - await this.sessionManager.setConnectionState(sessionId, false); - } - }, - async (error) => { - console.error(`Redis subscription error for session ${sessionId}:`, error); - await this.sessionManager.setConnectionState(sessionId, false); - } - ); - - return redisCleanup; - } - - /** - * Get the Redis channel name for a session (matches existing SSE pattern) - */ - private getRedisChannel(sessionId: string): string { - return `mcp:${sessionId}`; - } -} \ No newline at end of file diff --git a/src/services/redisTransport.ts b/src/services/redisTransport.ts new file mode 100644 index 0000000..2364088 --- /dev/null +++ b/src/services/redisTransport.ts @@ -0,0 +1,247 @@ +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import { redisClient } from "../redis.js"; +import { Transport, TransportSendOptions } from "@modelcontextprotocol/sdk/shared/transport.js"; +import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js"; +import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js"; +import { Server } from "@modelcontextprotocol/sdk/server/index.js"; + + +function getToServerChannel(sessionId: string): string { + return `mcp:shttp:toserver:${sessionId}`; +} + +function getToClientChannel(sessionId: string): string { + return `mcp:shttp:toclient:${sessionId}`; +} + +export async function isLive(sessionId: string): Promise { + // Check if the session is live by checking if the key exists in Redis + const numSubs = await redisClient.numsub(getToServerChannel(sessionId)); + console.log(`[isLive] Session ${sessionId}: Redis subscribers on ${getToServerChannel(sessionId)} = ${numSubs}`); + return numSubs > 0; +} + + +export class RedisTransport implements Transport { + private redisCleanup: (() => Promise) | undefined; + + onclose?: (() => void) | undefined; + onerror?: ((error: Error) => void) | undefined; + onmessage?: ((message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }) => void) | undefined; + + constructor( + private sendChannel: string, + private recvChannel: string, + private isLiveKey: string | undefined = undefined + ) { + this.sendChannel = sendChannel; + this.recvChannel = recvChannel; + this.isLiveKey = isLiveKey; + } + + + + async start(): Promise { + console.log(`[RedisTransport.start] Starting transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`); + if (this.redisCleanup) { + throw new Error(`Redis transport already started for channels ${this.sendChannel} and ${this.recvChannel}`); + } + + this.redisCleanup = await redisClient.createSubscription( + this.recvChannel, + (json) => { + console.log(`[RedisTransport] Received message on ${this.recvChannel}:`, json.substring(0, 100)); + const message = JSON.parse(json); + const extra = popExtra(message); + this.onmessage?.(message, extra) + }, + (error) => { + console.error( + `[Redis transport] Disconnecting due to error in Redis subscriber:`, + error, + ); + this.close() + }, + ); + console.log(`[RedisTransport.start] Successfully subscribed to ${this.recvChannel}`); + } + + send(message: JSONRPCMessage, options?: TransportSendOptions): Promise { + if (options) { + setOptions(message, options); + } + const messageStr = JSON.stringify(message); + console.log(`[RedisTransport.send] Publishing to ${this.sendChannel}:`, messageStr.substring(0, 100)); + return redisClient.publish(this.sendChannel, messageStr); + } + + + async cleanup(): Promise { + if (this.redisCleanup) { + await this.redisCleanup() + this.redisCleanup = undefined; + } + if (this.isLiveKey) { + await redisClient.del(this.isLiveKey); + } + } + + async close(): Promise { + console.log(`[RedisTransport.close] Closing transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`); + this.onclose?.(); + this.cleanup() + } +} + + +function setExtra(message: JSONRPCMessage, extra: { authInfo?: AuthInfo; } | undefined): void { + if (!extra) { + return; + } + if ("result" in message && typeof message.result === 'object') { + if (!message.result._meta) { + message.result._meta = {}; + } + message.result._meta.extra = extra; + } + if ("params" in message && typeof message.params === 'object') { + if (!message.params._meta) { + message.params._meta = {}; + } + message.params._meta.extra = extra; + } +} + +function setOptions(message: JSONRPCMessage, options: TransportSendOptions): void { + if ("result" in message && typeof message.result === 'object') { + if (!message.result._meta) { + message.result._meta = {}; + } + message.result._meta.options = options; + } + if ("params" in message && typeof message.params === 'object') { + if (!message.params._meta) { + message.params._meta = {}; + } + message.params._meta.options = options; + } +} + +function popExtra(message: JSONRPCMessage): { authInfo?: AuthInfo; } | undefined { + if ("params" in message && typeof message.params === 'object' && message.params._meta) { + const extra = message.params._meta.extra as { authInfo?: AuthInfo; } | undefined; + if (extra) { + delete message.params._meta.extra; + return extra; + } + } + if ("result" in message && typeof message.result === 'object' && message.result._meta) { + const extra = message.result._meta.extra as { authInfo?: AuthInfo; } | undefined; + if (extra) { + delete message.result._meta.extra; + return extra; + } + } + return undefined; +} + +function popOptions(message: JSONRPCMessage): TransportSendOptions | undefined { + if ("params" in message && typeof message.params === 'object' && message.params._meta) { + const options = message.params._meta.options as TransportSendOptions | undefined; + if (options) { + delete message.params._meta.options; + return options; + } + } + if ("result" in message && typeof message.result === 'object' && message.result._meta) { + const options = message.result._meta.options as TransportSendOptions | undefined; + if (options) { + delete message.result._meta.options; + return options; + } + } + return undefined; +} + + +function relayTransports(transport1: Transport, transport2: Transport): void { + console.log(`[relayTransports] Setting up relay between transports`); + transport1.onmessage = (message, extra) => { + console.log(`[relay] transport1 -> transport2:`, JSON.stringify(message).substring(0, 100)); + setExtra(message, extra); + const options = popOptions(message); + transport2.send(message, options) + }; + transport2.onmessage = (message, extra) => { + console.log(`[relay] transport2 -> transport1:`, JSON.stringify(message).substring(0, 100)); + setExtra(message, extra); + const options = popOptions(message); + transport1.send(message, options) + }; + + transport1.onerror = (error) => { + transport2.onerror?.(error); + }; + transport2.onerror = (error) => { + transport1.onerror?.(error); + }; + + transport1.onclose = () => { + transport2.close().catch(console.error); + }; + transport2.onclose = () => { + transport1.close().catch(console.error); + }; +} + + +export async function startServerListeningToRedis(server: Server, sessionId: string) { + console.log(`[startServerListeningToRedis] Starting background server for session ${sessionId}`); + const serverRedisTransport = createBackgroundTaskSideRedisTransport(sessionId) + await server.connect(serverRedisTransport) + console.log(`[startServerListeningToRedis] Background server connected for session ${sessionId}`); +} + +function createShttpHandlerSideRedisTransport(sessionId: string): RedisTransport { + return new RedisTransport( + getToServerChannel(sessionId), + getToClientChannel(sessionId), + ); +} + +function createBackgroundTaskSideRedisTransport(sessionId: string): RedisTransport { + return new RedisTransport( + getToClientChannel(sessionId), + getToServerChannel(sessionId), + ); +} + +export async function getFirstShttpTransport(sessionId: string): Promise { + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => sessionId, + enableJsonResponse: true, // Enable JSON response mode + }); + + const redisTransport = createShttpHandlerSideRedisTransport(sessionId); + + // When shttpTransport closes, so does the redisTransport + relayTransports(redisTransport, transport); + await redisTransport.start() + return transport; +} + +export async function getShttpTransport(sessionId: string): Promise { + // Giving undefined here and setting the sessionId means the + // transport wont try to create a new session. + const shttpTransport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined + }) + shttpTransport.sessionId = sessionId; + + const redisTransport = createShttpHandlerSideRedisTransport(sessionId); + + // When shttpTransport closes, so does the redisTransport + relayTransports(redisTransport, shttpTransport); + await redisTransport.start() + return shttpTransport; +} \ No newline at end of file diff --git a/src/services/sessionManager.test.ts b/src/services/sessionManager.test.ts deleted file mode 100644 index b9fac6b..0000000 --- a/src/services/sessionManager.test.ts +++ /dev/null @@ -1,160 +0,0 @@ -import { MockRedisClient } from "../redis.js"; -import { SessionManager } from "./sessionManager.js"; - -describe("SessionManager", () => { - let mockRedis: MockRedisClient; - let sessionManager: SessionManager; - - beforeEach(() => { - mockRedis = new MockRedisClient(); - sessionManager = new SessionManager(mockRedis); - }); - - afterEach(() => { - mockRedis.clear(); - }); - - describe("createSession", () => { - it("should create a new session with metadata", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - - const metadata = await sessionManager.getSessionMetadata(sessionId); - expect(metadata).not.toBeNull(); - expect(metadata!.sessionId).toBe(sessionId); - expect(metadata!.clientId).toBe(clientId); - expect(metadata!.createdAt).toBeGreaterThan(0); - expect(metadata!.lastActivity).toBeGreaterThan(0); - }); - - it("should initialize connection state as disconnected", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - - const isConnected = await sessionManager.isConnected(sessionId); - expect(isConnected).toBe(false); - }); - }); - - describe("refreshSession", () => { - it("should refresh an existing session", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - const originalMetadata = await sessionManager.getSessionMetadata(sessionId); - - // Wait a bit to ensure timestamp difference - await new Promise(resolve => setTimeout(resolve, 10)); - - const refreshed = await sessionManager.refreshSession(sessionId); - expect(refreshed).toBe(true); - - const newMetadata = await sessionManager.getSessionMetadata(sessionId); - expect(newMetadata!.lastActivity).toBeGreaterThan(originalMetadata!.lastActivity); - }); - - it("should return false for non-existent session", async () => { - const refreshed = await sessionManager.refreshSession("non-existent"); - expect(refreshed).toBe(false); - }); - }); - - describe("deleteSession", () => { - it("should delete a session and all associated data", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - await sessionManager.bufferMessage(sessionId, JSON.stringify({ test: "message" })); - - // Verify session exists - expect(await sessionManager.sessionExists(sessionId)).toBe(true); - - await sessionManager.deleteSession(sessionId); - - // Verify session is gone - expect(await sessionManager.sessionExists(sessionId)).toBe(false); - const messages = await sessionManager.getBufferedMessages(sessionId); - expect(messages).toHaveLength(0); - }); - }); - - describe("connection state management", () => { - it("should track connection state", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - - // Initially disconnected - expect(await sessionManager.isConnected(sessionId)).toBe(false); - - // Set to connected - await sessionManager.setConnectionState(sessionId, true); - expect(await sessionManager.isConnected(sessionId)).toBe(true); - - // Set back to disconnected - await sessionManager.setConnectionState(sessionId, false); - expect(await sessionManager.isConnected(sessionId)).toBe(false); - }); - }); - - describe("message buffering", () => { - it("should buffer and retrieve messages", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - - const message1 = JSON.stringify({ type: "request", id: 1 }); - const message2 = JSON.stringify({ type: "request", id: 2 }); - - await sessionManager.bufferMessage(sessionId, message1); - await sessionManager.bufferMessage(sessionId, message2); - - const bufferedMessages = await sessionManager.getBufferedMessages(sessionId); - expect(bufferedMessages).toHaveLength(2); - // Messages are returned in reverse order due to lpush - expect(bufferedMessages[0]).toBe(message2); - expect(bufferedMessages[1]).toBe(message1); - }); - - it("should clear buffered messages", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - await sessionManager.bufferMessage(sessionId, JSON.stringify({ test: "message" })); - - // Verify message is buffered - const messages = await sessionManager.getBufferedMessages(sessionId); - expect(messages).toHaveLength(1); - - // Clear messages - await sessionManager.clearBufferedMessages(sessionId); - - // Verify messages are cleared - const clearedMessages = await sessionManager.getBufferedMessages(sessionId); - expect(clearedMessages).toHaveLength(0); - }); - }); - - describe("sessionExists", () => { - it("should return true for existing session", async () => { - const sessionId = "test-session-123"; - const clientId = "test-client"; - - await sessionManager.createSession(sessionId, clientId); - expect(await sessionManager.sessionExists(sessionId)).toBe(true); - }); - - it("should return false for non-existent session", async () => { - expect(await sessionManager.sessionExists("non-existent")).toBe(false); - }); - }); -}); \ No newline at end of file diff --git a/src/services/sessionManager.ts b/src/services/sessionManager.ts deleted file mode 100644 index d8a9764..0000000 --- a/src/services/sessionManager.ts +++ /dev/null @@ -1,157 +0,0 @@ -import { RedisClient } from "../redis.js"; - -export interface SessionMetadata { - sessionId: string; - clientId: string; - createdAt: number; - lastActivity: number; -} - -/** - * Redis-based session manager for horizontal scaling - */ -export class SessionManager { - private static SESSION_TTL = 5 * 60; // 5 minutes in seconds - - constructor(private redisClient: RedisClient) {} - - // Redis key generators - private sessionMetadataKey(sessionId: string): string { - return `session:${sessionId}:metadata`; - } - - private sessionMessagesKey(sessionId: string): string { - return `session:${sessionId}:messages`; - } - - private sessionConnectionKey(sessionId: string): string { - return `session:${sessionId}:connection`; - } - - /** - * Create a new session in Redis with TTL - */ - async createSession(sessionId: string, clientId: string): Promise { - const metadata: SessionMetadata = { - sessionId, - clientId, - createdAt: Date.now(), - lastActivity: Date.now() - }; - - // Store session metadata with TTL - await this.redisClient.set( - this.sessionMetadataKey(sessionId), - JSON.stringify(metadata), - { EX: SessionManager.SESSION_TTL } - ); - - // Initialize empty message buffer - await this.redisClient.del(this.sessionMessagesKey(sessionId)); - - // Mark connection as disconnected initially - await this.redisClient.set( - this.sessionConnectionKey(sessionId), - 'disconnected', - { EX: SessionManager.SESSION_TTL } - ); - - console.log(`Session created: ${sessionId} for client: ${clientId}`); - } - - /** - * Refresh session TTL and update last activity - */ - async refreshSession(sessionId: string): Promise { - const metadata = await this.getSessionMetadata(sessionId); - if (!metadata) { - return false; - } - - // Update last activity and refresh TTL - metadata.lastActivity = Date.now(); - await this.redisClient.set( - this.sessionMetadataKey(sessionId), - JSON.stringify(metadata), - { EX: SessionManager.SESSION_TTL } - ); - - // Refresh other keys too - await this.redisClient.expire(this.sessionMessagesKey(sessionId), SessionManager.SESSION_TTL); - await this.redisClient.expire(this.sessionConnectionKey(sessionId), SessionManager.SESSION_TTL); - - console.log(`Session refreshed: ${sessionId}`); - return true; - } - - /** - * Delete session and all associated data - */ - async deleteSession(sessionId: string): Promise { - await this.redisClient.del(this.sessionMetadataKey(sessionId)); - await this.redisClient.del(this.sessionMessagesKey(sessionId)); - await this.redisClient.del(this.sessionConnectionKey(sessionId)); - - console.log(`Session deleted: ${sessionId}`); - } - - /** - * Get session metadata - */ - async getSessionMetadata(sessionId: string): Promise { - const data = await this.redisClient.get(this.sessionMetadataKey(sessionId)); - return data ? JSON.parse(data) : null; - } - - /** - * Check if session exists - */ - async sessionExists(sessionId: string): Promise { - const metadata = await this.getSessionMetadata(sessionId); - return metadata !== null; - } - - /** - * Mark SSE connection as connected/disconnected - */ - async setConnectionState(sessionId: string, connected: boolean): Promise { - await this.redisClient.set( - this.sessionConnectionKey(sessionId), - connected ? 'connected' : 'disconnected', - { EX: SessionManager.SESSION_TTL } - ); - - console.log(`Session ${sessionId} connection state: ${connected ? 'connected' : 'disconnected'}`); - } - - /** - * Check if SSE connection is active - */ - async isConnected(sessionId: string): Promise { - const state = await this.redisClient.get(this.sessionConnectionKey(sessionId)); - return state === 'connected'; - } - - /** - * Get buffered messages for a session - */ - async getBufferedMessages(sessionId: string): Promise { - return await this.redisClient.lrange(this.sessionMessagesKey(sessionId), 0, -1); - } - - /** - * Add a message to the buffer - */ - async bufferMessage(sessionId: string, message: string): Promise { - await this.redisClient.lpush(this.sessionMessagesKey(sessionId), message); - // Set TTL on the messages list - await this.redisClient.expire(this.sessionMessagesKey(sessionId), SessionManager.SESSION_TTL); - } - - /** - * Clear buffered messages - */ - async clearBufferedMessages(sessionId: string): Promise { - await this.redisClient.del(this.sessionMessagesKey(sessionId)); - } -} \ No newline at end of file