modelcontextprotocol/example-remote-server
Watch
1
0
Fork
You've already forked example-remote-server
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add comprehensive logging to debug Redis transport issues

Browse source 
- 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 <noreply@anthropic.com>
This commit is contained in:
Jerome 2025-06-30 20:29:41 +01:00
commit 8579d446aa
No known key found for this signature in database
13 changed files with 584 additions and 771 deletions
Download patch file Download diff file Expand all files Collapse all files

18
.vscode/launch.json vendored Normal file
View file

@ -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": ["<node_internals>/**"],
"env": {
"NODE_ENV": "development"
}
}
]
}

57
src/handlers/common.ts Normal file
View file

@ -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(),
);
}

344
src/handlers/mcp.ts
View file

@ -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();
}

115
src/handlers/shttp.ts Normal file
View file

@ -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 ===');
}
}

102
src/handlers/sse.ts Normal file
View file

@ -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();
}

17
src/index.ts
View file

@ -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

20
src/redis.ts
View file

@ -15,6 +15,8 @@ export interface RedisClient {
connect(): Promise<void>;
on(event: string, callback: (error: Error) => void): void;
options?: { url: string };
exists(key: string): Promise<boolean>;
numsub(key: string): Promise<number>;
/**
* 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<number> {
const subs = await this.redis.pubSubNumSub(key);
return subs[key] || 0;
}
async get(key: string): Promise<string | null> {
return await this.redis.get(key);
}
@ -110,6 +117,11 @@ export class RedisClientImpl implements RedisClient {
async publish(channel: string, message: string): Promise<void> {
await this.redis.publish(channel, message);
}
async exists(key: string): Promise<boolean> {
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<boolean> {
return this.store.has(key) || this.lists.has(key);
}
async numsub(key: string): Promise<number> {
return (this.subscribers.get(key) || []).length;
}
clear() {
this.store.clear();
this.lists.clear();

4
src/services/auth.ts
View file

@ -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");
}

12
src/services/mcp.ts
View file

@ -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<string> = new Set();
// Set up update interval for subscribed resources

102
src/services/messageDelivery.ts
View file

@ -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<void> {
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<void> {
// 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<void>> {
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}`;
}
}

247
src/services/redisTransport.ts Normal file
View file

@ -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<boolean> {
// 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<void>) | 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<void> {
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<void> {
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<void> {
if (this.redisCleanup) {
await this.redisCleanup()
this.redisCleanup = undefined;
}
if (this.isLiveKey) {
await redisClient.del(this.isLiveKey);
}
}
async close(): Promise<void> {
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<StreamableHTTPServerTransport> {
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<StreamableHTTPServerTransport> {
// 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;
}

160
src/services/sessionManager.test.ts
View file

@ -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);
});
});
});

157
src/services/sessionManager.ts
View file

@ -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<void> {
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<boolean> {
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<void> {
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<SessionMetadata | null> {
const data = await this.redisClient.get(this.sessionMetadataKey(sessionId));
return data ? JSON.parse(data) : null;
}
/**
* Check if session exists
*/
async sessionExists(sessionId: string): Promise<boolean> {
const metadata = await this.getSessionMetadata(sessionId);
return metadata !== null;
}
/**
* Mark SSE connection as connected/disconnected
*/
async setConnectionState(sessionId: string, connected: boolean): Promise<void> {
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<boolean> {
const state = await this.redisClient.get(this.sessionConnectionKey(sessionId));
return state === 'connected';
}
/**
* Get buffered messages for a session
*/
async getBufferedMessages(sessionId: string): Promise<string[]> {
return await this.redisClient.lrange(this.sessionMessagesKey(sessionId), 0, -1);
}
/**
* Add a message to the buffer
*/
async bufferMessage(sessionId: string, message: string): Promise<void> {
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<void> {
await this.redisClient.del(this.sessionMessagesKey(sessionId));
}
}