diff --git a/docs/streamable-http-design.md b/docs/streamable-http-design.md new file mode 100644 index 0000000..2d31081 --- /dev/null +++ b/docs/streamable-http-design.md @@ -0,0 +1,444 @@ +# Design Document: Implementing Streamable HTTP Transport for Example Remote Server + +## Research Summary + +### Current SSE Transport Architecture + +The example remote server currently uses the following architecture: + +1. **SSE Endpoint**: `/sse` - Creates SSE connection using `SSEServerTransport` +2. **Message Endpoint**: `/message` - Receives POST requests and forwards them via Redis +3. **Redis Integration**: Messages are published/subscribed through Redis channels using session IDs +4. **Auth**: Uses `requireBearerAuth` middleware with `EverythingAuthProvider` +5. **Session Management**: Each SSE connection gets a unique session ID used as Redis channel key + +**Key Files:** +- `/src/index.ts:91` - SSE endpoint with auth and headers +- `/src/handlers/mcp.ts:55-118` - SSE connection handler with Redis integration +- `/src/handlers/mcp.ts:120-144` - Message POST handler + +### Streamable HTTP Transport Specification (2025-03-26) + +The new Streamable HTTP transport replaces the old HTTP+SSE approach with a single endpoint that supports: + +1. **Single Endpoint**: One URL that handles GET, POST, and DELETE methods +2. **POST Requests**: Send JSON-RPC messages, can return either JSON responses or SSE streams +3. **GET Requests**: Open SSE streams for server-to-client messages +4. **Session Management**: Optional session IDs in `Mcp-Session-Id` headers +5. **Resumability**: Optional event storage with `Last-Event-ID` support +6. **Auth Integration**: Same authentication patterns as SSE + +**Key Specification Requirements:** +- Accept header must include both `application/json` and `text/event-stream` +- Session ID management via `Mcp-Session-Id` headers +- 202 Accepted for notifications/responses only +- SSE streams or JSON responses for requests +- Security: Origin validation, localhost binding, proper auth + +### TypeScript SDK Implementation + +The SDK provides `StreamableHTTPServerTransport` with: + +1. **Two Modes**: + - **Stateful**: Session ID generator provided, maintains sessions in memory + - **Stateless**: Session ID generator undefined, no session state + +2. **Key Features**: + - Built-in session validation + - Event store support for resumability + - Automatic response correlation + - Auth info threading via `req.auth` + +3. **Integration Patterns**: + - **Stateful**: Store transports by session ID, reuse across requests + - **Stateless**: New transport per request, immediate cleanup + - **Auth**: Same bearer auth middleware as SSE + +## Implementation Plan + +### 1. New Streamable HTTP Endpoint + +Add `/mcp` endpoint that handles GET, POST, DELETE methods: + +```typescript +// In src/index.ts +app.get("/mcp", cors(corsOptions), bearerAuth, authContext, handleStreamableHTTP); +app.post("/mcp", cors(corsOptions), bearerAuth, authContext, handleStreamableHTTP); +app.delete("/mcp", cors(corsOptions), bearerAuth, authContext, handleStreamableHTTP); +``` + +### 2. Handler Implementation + +Create new handler in `/src/handlers/mcp.ts`: + +```typescript +export async function handleStreamableHTTP(req: Request, res: Response) { + // Use same Redis-based architecture as SSE transport + // but with StreamableHTTPServerTransport instead of SSEServerTransport +} +``` + +### 3. Transport Integration Strategy (Horizontally Scalable) + +**Redis-based Session Management (Required for Horizontal Scaling)** +- Store session state in Redis, not in-memory +- Any server instance can handle any request for any session +- Session lifecycle independent of SSE connection lifecycle +- Message buffering in Redis when SSE connection is down +- Session TTL of 5 minutes to prevent Redis bloat + +### 4. Redis Integration + +Maintain current Redis architecture: +- Use session ID as Redis channel key +- Same message publishing/subscribing pattern +- Same MCP server creation logic +- Transport acts as bridge to Redis like current SSE implementation + +### 5. Auth Integration + +Use identical auth setup as SSE: +- Same `bearerAuth` middleware +- Same `authContext` middleware +- Same `EverythingAuthProvider` +- Auth info flows through `req.auth` to transport + +### 6. Backwards Compatibility + +Keep existing `/sse` and `/message` endpoints: +- Maintain current SSE transport for existing clients +- Add new `/mcp` endpoint alongside +- Both transports share same Redis infrastructure +- Same auth provider serves both + +## Key Differences from Current SSE Implementation + +1. **Single Endpoint**: `/mcp` handles all HTTP methods vs separate `/sse` + `/message` +2. **Transport Class**: `StreamableHTTPServerTransport` vs `SSEServerTransport` +3. **Session Headers**: `Mcp-Session-Id` headers vs URL session ID +4. **Request Handling**: Transport handles HTTP details vs manual SSE headers +5. **Response Correlation**: Built into transport vs manual request tracking + +## Benefits of This Approach + +1. **Spec Compliance**: Follows 2025-03-26 MCP specification exactly +2. **Minimal Changes**: Reuses existing Redis infrastructure and auth +3. **Feature Parity**: Same functionality as current SSE transport +4. **Future Proof**: Can add resumability with event store later +5. **Clean Integration**: Same auth patterns and middleware stack + +## Implementation Steps + +1. **Add Dependencies**: `StreamableHTTPServerTransport` from SDK +2. **Create Redis Session Management**: Implement `SessionManager` and `MessageDelivery` classes +3. **Create Handler**: New streamable HTTP handler function with Redis integration +4. **Add Routes**: New `/mcp` endpoint with all HTTP methods +5. **Session Management**: Redis-based session storage with TTL +6. **Message Buffering**: Redis-based message buffering for disconnected clients +7. **Testing**: Verify auth, Redis integration, horizontal scaling, and MCP protocol compliance +8. **Documentation**: Update README with new endpoint usage + +### Additional Implementation Considerations + +- **Redis Connection Management**: Ensure Redis connections are properly pooled and cleaned up +- **Error Handling**: Robust error handling for Redis operations and session timeouts +- **Monitoring**: Add logging for session creation, cleanup, and message buffering metrics +- **Performance**: Consider Redis memory usage and implement appropriate limits on message buffer size +- **Security**: Ensure session IDs are cryptographically secure and validate all session operations + +## Technical Details + +### Session Management Architecture (Redis-based) + +**Redis Data Structures for Horizontal Scaling:** + +```typescript +// Redis keys for session management +const SESSION_METADATA_KEY = (sessionId: string) => `session:${sessionId}:metadata`; +const SESSION_MESSAGES_KEY = (sessionId: string) => `session:${sessionId}:messages`; +const SESSION_CONNECTION_KEY = (sessionId: string) => `session:${sessionId}:connection`; + +// Session metadata structure +interface SessionMetadata { + sessionId: string; + clientId: string; + createdAt: number; + lastActivity: number; +} + +// Session lifecycle management +class SessionManager { + private static SESSION_TTL = 5 * 60; // 5 minutes in seconds + + static async createSession(sessionId: string, clientId: string): Promise { + const metadata: SessionMetadata = { + sessionId, + clientId, + createdAt: Date.now(), + lastActivity: Date.now() + }; + + // Store session metadata with TTL + await redisClient.set( + SESSION_METADATA_KEY(sessionId), + JSON.stringify(metadata), + { EX: this.SESSION_TTL } + ); + + // Initialize empty message buffer + await redisClient.del(SESSION_MESSAGES_KEY(sessionId)); + + // Mark connection as disconnected initially + await redisClient.set(SESSION_CONNECTION_KEY(sessionId), 'disconnected', { EX: this.SESSION_TTL }); + } + + static 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 redisClient.set( + SESSION_METADATA_KEY(sessionId), + JSON.stringify(metadata), + { EX: this.SESSION_TTL } + ); + + // Refresh other keys too + await redisClient.expire(SESSION_MESSAGES_KEY(sessionId), this.SESSION_TTL); + await redisClient.expire(SESSION_CONNECTION_KEY(sessionId), this.SESSION_TTL); + + return true; + } + + static async deleteSession(sessionId: string): Promise { + await redisClient.del(SESSION_METADATA_KEY(sessionId)); + await redisClient.del(SESSION_MESSAGES_KEY(sessionId)); + await redisClient.del(SESSION_CONNECTION_KEY(sessionId)); + } + + static async getSessionMetadata(sessionId: string): Promise { + const data = await redisClient.get(SESSION_METADATA_KEY(sessionId)); + return data ? JSON.parse(data) : null; + } + + // Mark SSE connection as connected/disconnected + static async setConnectionState(sessionId: string, connected: boolean): Promise { + await redisClient.set( + SESSION_CONNECTION_KEY(sessionId), + connected ? 'connected' : 'disconnected', + { EX: this.SESSION_TTL } + ); + } + + static async isConnected(sessionId: string): Promise { + const state = await redisClient.get(SESSION_CONNECTION_KEY(sessionId)); + return state === 'connected'; + } +} +``` + +### Redis Integration Pattern with Message Buffering + +The implementation extends the current Redis pattern to support message buffering: + +```typescript +// Message delivery with buffering support +class MessageDelivery { + static async deliverMessage(sessionId: string, message: JSONRPCMessage): Promise { + const isConnected = await SessionManager.isConnected(sessionId); + + if (isConnected) { + // Direct delivery via existing Redis pub/sub + const redisChannel = `mcp:${sessionId}`; + await redisClient.publish(redisChannel, JSON.stringify(message)); + } else { + // Buffer the message for later delivery + await redisClient.lpush( + SESSION_MESSAGES_KEY(sessionId), + JSON.stringify(message) + ); + // Set TTL on the messages list + await redisClient.expire(SESSION_MESSAGES_KEY(sessionId), SessionManager.SESSION_TTL); + } + } + + static async deliverBufferedMessages(sessionId: string, transport: StreamableHTTPServerTransport): Promise { + // Get all buffered messages + const bufferedMessages = await redisClient.lrange(SESSION_MESSAGES_KEY(sessionId), 0, -1); + + // Deliver buffered messages in order (reverse because lpush) + for (let i = bufferedMessages.length - 1; i >= 0; i--) { + const message = JSON.parse(bufferedMessages[i]); + await transport.send(message); + } + + // Clear the buffer after delivery + await redisClient.del(SESSION_MESSAGES_KEY(sessionId)); + } +} + +// Enhanced Redis subscription for SSE connections +const setupRedisSubscription = async (sessionId: string, transport: StreamableHTTPServerTransport) => { + const redisChannel = `mcp:${sessionId}`; + + const redisCleanup = await redisClient.createSubscription( + redisChannel, + async (message) => { + const jsonMessage = JSON.parse(message); + try { + await transport.send(jsonMessage); + } catch (error) { + console.error(`Failed to send message on transport for session ${sessionId}:`, error); + // Mark connection as disconnected so future messages get buffered + await SessionManager.setConnectionState(sessionId, false); + } + }, + async (error) => { + console.error('Redis subscription error:', error); + await SessionManager.setConnectionState(sessionId, false); + } + ); + + return redisCleanup; +}; +``` + +### Handler Implementation Flow + +The new streamable HTTP handler integrates with the Redis-based session management: + +```typescript +export async function handleStreamableHTTP(req: Request, res: Response) { + const method = req.method; + const sessionId = req.headers['mcp-session-id'] as string | undefined; + + if (method === 'POST') { + // Handle POST requests (initialization or message sending) + + if (isInitializeRequest(req.body) && !sessionId) { + // New session initialization + const newSessionId = randomUUID(); + const authInfo = req.auth; + + // Create session in Redis + await SessionManager.createSession(newSessionId, authInfo?.clientId || 'unknown'); + + // Create transport with Redis-based session management + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => newSessionId, + // Custom implementation - don't store transport in memory + }); + + const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer(); + + // Set up Redis subscription for this session but don't store transport globally + // Instead, rely on Redis for all message routing + + await mcpServer.connect(transport); + await transport.handleRequest(req, res, req.body); + + } 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; + } + + // Create ephemeral transport for this request + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => sessionId, // Use existing session ID + }); + + const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer(); + await mcpServer.connect(transport); + await transport.handleRequest(req, res, req.body); + + // Clean up after request completes + res.on('close', mcpCleanup); + } + + } else if (method === 'GET') { + // Handle SSE stream requests + + if (!sessionId) { + res.writeHead(400).end('Session ID required'); + return; + } + + const sessionValid = await SessionManager.refreshSession(sessionId); + if (!sessionValid) { + res.writeHead(404).end('Session not found'); + return; + } + + // 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 setupRedisSubscription(sessionId, transport); + + // Handle connection cleanup + res.on('close', async () => { + await SessionManager.setConnectionState(sessionId, false); + redisCleanup(); + }); + + await transport.handleRequest(req, res); + + } else if (method === 'DELETE') { + // Handle session deletion + + if (!sessionId) { + res.writeHead(400).end('Session ID required'); + return; + } + + // Delete session from Redis + await SessionManager.deleteSession(sessionId); + res.writeHead(200).end(); + } +} +``` + +### Auth Information Flow + +Auth information flows through the middleware stack: + +```typescript +// Auth middleware adds req.auth +const authInfo: AuthInfo = req.auth; + +// Transport receives auth info +await transport.handleRequest(req, res); + +// Auth info is available in MCP server handlers +server.tool('example', 'description', schema, async (params, { authInfo }) => { + // authInfo contains token, clientId, scopes, etc. +}); +``` + +## Conclusion + +This design provides horizontally scalable streamable HTTP support by using Redis for all session state management and message buffering. Key advantages: + +1. **Horizontal Scalability**: Any server instance can handle any request for any session +2. **Resilient Connection Handling**: SSE disconnects don't end sessions; messages are buffered +3. **Automatic Cleanup**: 5-minute session TTL prevents Redis bloat +4. **Backwards Compatibility**: Existing `/sse` and `/message` endpoints remain unchanged +5. **Spec Compliance**: Follows 2025-03-26 MCP specification exactly + +The implementation is more complex than a single-instance approach but essential for production deployment in a horizontally scaled environment. The Redis-based architecture ensures sessions persist across server instances and SSE connection interruptions. \ No newline at end of file diff --git a/src/handlers/mcp.ts b/src/handlers/mcp.ts index 1a70833..8a4b531 100644 --- a/src/handlers/mcp.ts +++ b/src/handlers/mcp.ts @@ -1,11 +1,16 @@ 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"; @@ -142,3 +147,198 @@ export async function handleMessage(req: Request, res: Response) { 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/index.ts b/src/index.ts index 4fc70ab..e2e6dd0 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,7 +3,7 @@ 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, authContext } from "./handlers/mcp.js"; +import { handleMessage, handleSSEConnection, handleStreamableHTTP, authContext } from "./handlers/mcp.js"; import { handleFakeAuthorizeRedirect, handleFakeAuthorize } from "./handlers/fakeauth.js"; import { redisClient } from "./redis.js"; import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js"; @@ -87,10 +87,15 @@ const options: AuthRouterOptions = { app.use(mcpAuthRouter(options)); const bearerAuth = requireBearerAuth(options); -// MCP routes +// MCP routes (legacy SSE transport) app.get("/sse", cors(corsOptions), bearerAuth, authContext, sseHeaders, handleSSEConnection); app.post("/message", cors(corsOptions), bearerAuth, authContext, sensitiveDataHeaders, handleMessage); +// MCP routes (new streamable HTTP transport) +app.get("/mcp", cors(corsOptions), bearerAuth, authContext, handleStreamableHTTP); +app.post("/mcp", cors(corsOptions), bearerAuth, authContext, handleStreamableHTTP); +app.delete("/mcp", cors(corsOptions), bearerAuth, authContext, handleStreamableHTTP); + // Upstream auth routes app.get("/fakeupstreamauth/authorize", cors(corsOptions), handleFakeAuthorize); app.get("/fakeupstreamauth/callback", cors(corsOptions), handleFakeAuthorizeRedirect); diff --git a/src/redis.ts b/src/redis.ts index 7bbdc7e..dda7a27 100644 --- a/src/redis.ts +++ b/src/redis.ts @@ -8,6 +8,10 @@ export interface RedisClient { get(key: string): Promise; set(key: string, value: string, options?: SetOptions): Promise; getDel(key: string): Promise; + del(key: string): Promise; + expire(key: string, seconds: number): Promise; + lpush(key: string, ...values: string[]): Promise; + lrange(key: string, start: number, stop: number): Promise; connect(): Promise; on(event: string, callback: (error: Error) => void): void; options?: { url: string }; @@ -60,6 +64,22 @@ export class RedisClientImpl implements RedisClient { ); } + async del(key: string): Promise { + return await this.redis.del(key); + } + + async expire(key: string, seconds: number): Promise { + return await this.redis.expire(key, seconds); + } + + async lpush(key: string, ...values: string[]): Promise { + return await this.redis.lPush(key, values); + } + + async lrange(key: string, start: number, stop: number): Promise { + return await this.redis.lRange(key, start, stop); + } + async connect(): Promise { await this.redis.connect(); } @@ -103,6 +123,7 @@ export function setRedisClient(client: RedisClient) { export class MockRedisClient implements RedisClient { options = { url: "redis://localhost:6379" }; private store = new Map(); + private lists = new Map(); private subscribers = new Map void)[]>(); private errorCallbacks = new Map void)[]>(); @@ -125,6 +146,39 @@ export class MockRedisClient implements RedisClient { return oldValue; } + async del(key: string): Promise { + let deleted = 0; + if (this.store.has(key)) { + this.store.delete(key); + deleted++; + } + if (this.lists.has(key)) { + this.lists.delete(key); + deleted++; + } + return deleted; + } + + async expire(key: string, _seconds: number): Promise { + // Mock implementation - just return true if key exists + return this.store.has(key) || this.lists.has(key); + } + + async lpush(key: string, ...values: string[]): Promise { + const list = this.lists.get(key) || []; + list.unshift(...values); + this.lists.set(key, list); + return list.length; + } + + async lrange(key: string, start: number, stop: number): Promise { + const list = this.lists.get(key) || []; + if (stop === -1) { + return list.slice(start); + } + return list.slice(start, stop + 1); + } + async connect(): Promise { // No-op in mock } @@ -185,6 +239,7 @@ export class MockRedisClient implements RedisClient { clear() { this.store.clear(); + this.lists.clear(); this.subscribers.clear(); this.errorCallbacks.clear(); } diff --git a/src/services/messageDelivery.ts b/src/services/messageDelivery.ts new file mode 100644 index 0000000..1fe6cba --- /dev/null +++ b/src/services/messageDelivery.ts @@ -0,0 +1,102 @@ +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/sessionManager.test.ts b/src/services/sessionManager.test.ts new file mode 100644 index 0000000..b9fac6b --- /dev/null +++ b/src/services/sessionManager.test.ts @@ -0,0 +1,160 @@ +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 new file mode 100644 index 0000000..d8a9764 --- /dev/null +++ b/src/services/sessionManager.ts @@ -0,0 +1,157 @@ +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