Implement streamable HTTP transport with Redis-based session management

- Add new /mcp endpoint supporting GET, POST, DELETE methods
- Implement SessionManager for Redis-based session state management
- Add MessageDelivery for message buffering when SSE disconnected
- Support horizontal scaling via Redis session storage
- Maintain backwards compatibility with existing /sse and /message endpoints
- Add comprehensive tests for SessionManager functionality
- Session TTL of 5 minutes prevents Redis bloat
- Message buffering for resilient connection handling

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

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jerome 2025-06-05 12:24:04 +01:00
commit 2b673a28c2
No known key found for this signature in database
6 changed files with 1125 additions and 2 deletions

View file

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

View file

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

View file

@ -8,6 +8,10 @@ export interface RedisClient {
get(key: string): Promise<string | null>;
set(key: string, value: string, options?: SetOptions): Promise<string | null>;
getDel(key: string): Promise<string | null>;
del(key: string): Promise<number>;
expire(key: string, seconds: number): Promise<boolean>;
lpush(key: string, ...values: string[]): Promise<number>;
lrange(key: string, start: number, stop: number): Promise<string[]>;
connect(): Promise<void>;
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<number> {
return await this.redis.del(key);
}
async expire(key: string, seconds: number): Promise<boolean> {
return await this.redis.expire(key, seconds);
}
async lpush(key: string, ...values: string[]): Promise<number> {
return await this.redis.lPush(key, values);
}
async lrange(key: string, start: number, stop: number): Promise<string[]> {
return await this.redis.lRange(key, start, stop);
}
async connect(): Promise<void> {
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<string, string>();
private lists = new Map<string, string[]>();
private subscribers = new Map<string, ((message: string) => void)[]>();
private errorCallbacks = new Map<string, ((error: Error) => void)[]>();
@ -125,6 +146,39 @@ export class MockRedisClient implements RedisClient {
return oldValue;
}
async del(key: string): Promise<number> {
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<boolean> {
// Mock implementation - just return true if key exists
return this.store.has(key) || this.lists.has(key);
}
async lpush(key: string, ...values: string[]): Promise<number> {
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<string[]> {
const list = this.lists.get(key) || [];
if (stop === -1) {
return list.slice(start);
}
return list.slice(start, stop + 1);
}
async connect(): Promise<void> {
// 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();
}

View file

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

View file

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

View file

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