From c2119ed9f97da48d71b2e50aeca4acd1f0eb7dd2 Mon Sep 17 00:00:00 2001 From: Jerome Date: Tue, 1 Jul 2025 15:10:55 +0100 Subject: [PATCH] WIP shttp changes --- src/handlers/shttp.ts | 46 +++++++++++------------- src/index.ts | 15 ++++++-- src/services/mcp.ts | 19 +++------- src/services/redisTransport.ts | 66 +++++++++++++++++++++++++--------- 4 files changed, 88 insertions(+), 58 deletions(-) diff --git a/src/handlers/shttp.ts b/src/handlers/shttp.ts index 79eba9b..84940b5 100644 --- a/src/handlers/shttp.ts +++ b/src/handlers/shttp.ts @@ -1,11 +1,10 @@ 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 { getFirstShttpTransport, getShttpTransport, isLive, RedisTransport, 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" { @@ -22,7 +21,8 @@ export async function handleStreamableHTTP(req: Request, res: Response) { 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; + let shttpTransport: StreamableHTTPServerTransport | undefined = undefined; + let redisTransport: RedisTransport | undefined = undefined; try { // Check for existing session ID const sessionId = req.headers['mcp-session-id'] as string | undefined; @@ -31,7 +31,11 @@ export async function handleStreamableHTTP(req: Request, res: Response) { if (sessionId && await isLive(sessionId)) { console.log('Session is live, reusing existing transport for session:', sessionId); // Reuse existing transport - transport = await getShttpTransport(sessionId) + // const { shttpTransport, redisTransport } = await getShttpTransport(sessionId); + // destructuring to get both transports + ({ shttpTransport, redisTransport } = await getShttpTransport(sessionId)); + + console.log('Created transport for session:', 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'); @@ -45,20 +49,9 @@ export async function handleStreamableHTTP(req: Request, res: Response) { startServerListeningToRedis(server.server, sessionId) console.log('Started server listening to Redis for session:', sessionId); - transport = await getFirstShttpTransport(sessionId); + ({ shttpTransport, redisTransport } = 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 + console.log('Transport object:', shttpTransport.constructor.name, 'sessionId:', shttpTransport.sessionId); } else { console.log('Invalid request - no session ID and not an initialization request'); console.log('Session ID present:', !!sessionId); @@ -79,9 +72,9 @@ export async function handleStreamableHTTP(req: Request, res: Response) { // 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('Transport object:', shttpTransport.constructor.name, 'sessionId:', shttpTransport.sessionId); + console.log('Request body:', JSON.stringify(req.body, null, 2)); + await shttpTransport.handleRequest(req, res, req.body); console.log('Request handled successfully'); } catch (error) { console.error('=== ERROR in handleStreamableHTTP ==='); @@ -104,12 +97,13 @@ export async function handleStreamableHTTP(req: Request, res: Response) { 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'); - // } + // Set up cleanup when response is complete + res.on('finish', async () => { + console.log('HTTP response finished, closing transport'); + await shttpTransport?.close(); + // await redisTransport?.close(); + console.log('Transport closed after response'); + }); console.log('=== handleStreamableHTTP END ==='); } } diff --git a/src/index.ts b/src/index.ts index 9d67291..26eb3f2 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,4 +1,4 @@ -import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js"; +import { BearerAuthMiddlewareOptions, 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"; @@ -9,6 +9,7 @@ import { handleFakeAuthorize, handleFakeAuthorizeRedirect } from "./handlers/fak import { handleStreamableHTTP } from "./handlers/shttp.js"; import { handleMessage, handleSSEConnection } from "./handlers/sse.js"; import { redisClient } from "./redis.js"; +import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js"; const app = express(); @@ -78,6 +79,8 @@ app.use(baseSecurityHeaders); // Enable CORS pre-flight requests app.options('*', cors(corsOptions)); + +const authProvider = new EverythingAuthProvider(); // Auth configuration const options: AuthRouterOptions = { provider: new EverythingAuthProvider(), @@ -89,8 +92,16 @@ const options: AuthRouterOptions = { } } }; + +const dearerAuthMiddlewareOptions: BearerAuthMiddlewareOptions = { + // verifyAccessToken(token: string): Promise; + verifier: { + verifyAccessToken: authProvider.verifyAccessToken.bind(authProvider), + } +} + app.use(mcpAuthRouter(options)); -const bearerAuth = requireBearerAuth(options); +const bearerAuth = requireBearerAuth(dearerAuthMiddlewareOptions); // MCP routes (legacy SSE transport) app.get("/sse", cors(corsOptions), bearerAuth, authContext, sseHeaders, handleSSEConnection); diff --git a/src/services/mcp.ts b/src/services/mcp.ts index 8d8711e..9551e6b 100644 --- a/src/services/mcp.ts +++ b/src/services/mcp.ts @@ -15,13 +15,16 @@ import { SetLevelRequestSchema, SubscribeRequestSchema, Tool, - ToolSchema, UnsubscribeRequestSchema, } from "@modelcontextprotocol/sdk/types.js"; import { z } from "zod"; import { zodToJsonSchema } from "zod-to-json-schema"; -type ToolInput = z.infer; +type ToolInput = { + type: "object"; + properties?: Record; + required?: string[]; +}; /* Input schemas for tools implemented in this server */ const EchoSchema = z.object({ @@ -109,18 +112,6 @@ export const createMcpServer = () => { } ); - // Add logging to track all incoming messages - const originalSetRequestHandler = server.setRequestHandler.bind(server); - server.setRequestHandler = function(schema: any, handler: any) { - const wrappedHandler = async (request: any) => { - console.log(`[MCP Server] Handling request: ${request.method}`); - const result = await handler(request); - console.log(`[MCP Server] Request ${request.method} completed`); - return result; - }; - return originalSetRequestHandler(schema, wrappedHandler); - }; - const subscriptions: Set = new Set(); // Set up update interval for subscribed resources diff --git a/src/services/redisTransport.ts b/src/services/redisTransport.ts index 2364088..297c059 100644 --- a/src/services/redisTransport.ts +++ b/src/services/redisTransport.ts @@ -5,6 +5,7 @@ import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js"; import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js"; import { Server } from "@modelcontextprotocol/sdk/server/index.js"; +let redisTransportCounter = 0; function getToServerChannel(sessionId: string): string { return `mcp:shttp:toserver:${sessionId}`; @@ -24,6 +25,7 @@ export async function isLive(sessionId: string): Promise { export class RedisTransport implements Transport { private redisCleanup: (() => Promise) | undefined; + private counter: number; onclose?: (() => void) | undefined; onerror?: ((error: Error) => void) | undefined; @@ -34,6 +36,7 @@ export class RedisTransport implements Transport { private recvChannel: string, private isLiveKey: string | undefined = undefined ) { + this.counter = redisTransportCounter++; this.sendChannel = sendChannel; this.recvChannel = recvChannel; this.isLiveKey = isLiveKey; @@ -42,18 +45,26 @@ export class RedisTransport implements Transport { async start(): Promise { - console.log(`[RedisTransport.start] Starting transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`); + console.log(`[RedisTransport.${this.counter}.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}`); } + // Log when onmessage is set + console.log(`[RedisTransport.${this.counter}.start] onmessage handler is ${this.onmessage ? 'SET' : 'NOT SET'}`); + this.redisCleanup = await redisClient.createSubscription( this.recvChannel, (json) => { - console.log(`[RedisTransport] Received message on ${this.recvChannel}:`, json.substring(0, 100)); + console.log(`[RedisTransport.${this.counter}] Received message on ${this.recvChannel}:`, json.substring(0, 100)); const message = JSON.parse(json); const extra = popExtra(message); - this.onmessage?.(message, extra) + if (this.onmessage) { + console.log(`[RedisTransport.${this.counter}] Calling onmessage handler for ${this.recvChannel}`); + this.onmessage(message, extra) + } else { + console.log(`[RedisTransport.${this.counter}] WARNING: No onmessage handler for ${this.recvChannel}!`); + } }, (error) => { console.error( @@ -63,23 +74,27 @@ export class RedisTransport implements Transport { this.close() }, ); - console.log(`[RedisTransport.start] Successfully subscribed to ${this.recvChannel}`); + console.log(`[RedisTransport.${this.counter}.start] Successfully subscribed to ${this.recvChannel}`); } - send(message: JSONRPCMessage, options?: TransportSendOptions): Promise { + async send(message: JSONRPCMessage, options?: TransportSendOptions): Promise { if (options) { setOptions(message, options); } const messageStr = JSON.stringify(message); - console.log(`[RedisTransport.send] Publishing to ${this.sendChannel}:`, messageStr.substring(0, 100)); - return redisClient.publish(this.sendChannel, messageStr); + console.log(`[RedisTransport.${this.counter}.send] Publishing to ${this.sendChannel}:`, messageStr.substring(0, 100)); + console.log(`[RedisTransport.${this.counter}.send] Full message:`, messageStr); + await redisClient.publish(this.sendChannel, messageStr); + console.log(`[RedisTransport.${this.counter}.send] Published successfully to ${this.sendChannel}`); } async cleanup(): Promise { if (this.redisCleanup) { + console.log(`[RedisTransport.${this.counter}.cleanup] Unsubscribing from ${this.recvChannel}`); await this.redisCleanup() this.redisCleanup = undefined; + console.log(`[RedisTransport.${this.counter}.cleanup] Successfully unsubscribed from ${this.recvChannel}`); } if (this.isLiveKey) { await redisClient.del(this.isLiveKey); @@ -87,9 +102,9 @@ export class RedisTransport implements Transport { } async close(): Promise { - console.log(`[RedisTransport.close] Closing transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`); + console.log(`[RedisTransport.${this.counter}.close] Closing transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`); this.onclose?.(); - this.cleanup() + await this.cleanup() } } @@ -186,11 +201,20 @@ function relayTransports(transport1: Transport, transport2: Transport): void { transport1.onerror?.(error); }; + // Prevent circular close calls + let closing = false; + transport1.onclose = () => { - transport2.close().catch(console.error); + if (!closing) { + closing = true; + transport2.close().catch(console.error); + } }; transport2.onclose = () => { - transport1.close().catch(console.error); + if (!closing) { + closing = true; + transport1.close().catch(console.error); + } }; } @@ -198,6 +222,15 @@ function relayTransports(transport1: Transport, transport2: Transport): void { export async function startServerListeningToRedis(server: Server, sessionId: string) { console.log(`[startServerListeningToRedis] Starting background server for session ${sessionId}`); const serverRedisTransport = createBackgroundTaskSideRedisTransport(sessionId) + + // Log all messages sent by the server + const originalSend = serverRedisTransport.send.bind(serverRedisTransport); + serverRedisTransport.send = async (message, options) => { + console.log(`[MCP Server -> Redis] Sending response:`, JSON.stringify(message).substring(0, 200)); + return originalSend(message, options); + }; + + // The server.connect() will call start() on the transport await server.connect(serverRedisTransport) console.log(`[startServerListeningToRedis] Background server connected for session ${sessionId}`); } @@ -216,7 +249,7 @@ function createBackgroundTaskSideRedisTransport(sessionId: string): RedisTranspo ); } -export async function getFirstShttpTransport(sessionId: string): Promise { +export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, redisTransport: RedisTransport}> { const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => sessionId, enableJsonResponse: true, // Enable JSON response mode @@ -227,14 +260,15 @@ export async function getFirstShttpTransport(sessionId: string): Promise { +export async function getShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, redisTransport: RedisTransport}> { // Giving undefined here and setting the sessionId means the // transport wont try to create a new session. const shttpTransport = new StreamableHTTPServerTransport({ - sessionIdGenerator: undefined + sessionIdGenerator: undefined, + enableJsonResponse: true, // Use JSON response mode for all requests }) shttpTransport.sessionId = sessionId; @@ -243,5 +277,5 @@ export async function getShttpTransport(sessionId: string): Promise