From de8c4a4ed2de4a2ef9fd96023a451b8f68a98c7d Mon Sep 17 00:00:00 2001 From: Jerome Date: Fri, 4 Jul 2025 17:23:54 +0100 Subject: [PATCH] Fix Redis transport refactoring with control message system MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix timing issue by awaiting startServerListeningToRedis - Add request-id based response channel routing - Implement control message system for proper server shutdown - Add comprehensive integration tests - Ensure response subscription is established before message sending - Fix MCP server cleanup integration to prevent async hanging 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- src/handlers/shttp.ts | 2 +- src/redis.ts | 20 +- src/services/auth.test.ts | 6 + src/services/auth.ts | 2 +- src/services/mcp.ts | 7 +- .../redisTransport.integration.test.ts | 256 ++++++++++++++++++ src/services/redisTransport.test.ts | 135 +++++++-- src/services/redisTransport.ts | 184 ++++++++++--- 8 files changed, 551 insertions(+), 61 deletions(-) create mode 100644 src/services/redisTransport.integration.test.ts diff --git a/src/handlers/shttp.ts b/src/handlers/shttp.ts index 894813c..53bd80a 100644 --- a/src/handlers/shttp.ts +++ b/src/handlers/shttp.ts @@ -44,7 +44,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) { const server = createMcpServer(); console.log('Created MCP server instance'); - startServerListeningToRedis(server.server, sessionId) + await startServerListeningToRedis(server, sessionId) console.log('Started server listening to Redis for session:', sessionId); ({ shttpTransport, cleanup } = await getFirstShttpTransport(sessionId)); diff --git a/src/redis.ts b/src/redis.ts index 5ad7d87..e45e85e 100644 --- a/src/redis.ts +++ b/src/redis.ts @@ -104,18 +104,34 @@ export class RedisClientImpl implements RedisClient { onMessage: (message: string) => void, onError: (error: Error) => void, ): Promise<() => Promise> { + console.log(`[RedisClient] Creating subscription for channel: ${channel}`); const subscriber = this.redis.duplicate(); - subscriber.on("error", onError); + subscriber.on("error", (error) => { + console.error(`[RedisClient] Subscriber error for channel ${channel}:`, error); + onError(error); + }); + + console.log(`[RedisClient] Connecting subscriber for channel: ${channel}`); await subscriber.connect(); - await subscriber.subscribe(channel, onMessage); + console.log(`[RedisClient] Connected subscriber for channel: ${channel}`); + + console.log(`[RedisClient] Subscribing to channel: ${channel}`); + await subscriber.subscribe(channel, (message) => { + console.log(`[RedisClient] Received message on channel ${channel}:`, message.substring(0, 100)); + onMessage(message); + }); + console.log(`[RedisClient] Successfully subscribed to channel: ${channel}`); return async () => { + console.log(`[RedisClient] Cleaning up subscription for channel: ${channel}`); await subscriber.disconnect(); }; } async publish(channel: string, message: string): Promise { + console.log(`[RedisClient] Publishing to channel: ${channel}, message length: ${message.length}`); await this.redis.publish(channel, message); + console.log(`[RedisClient] Successfully published to channel: ${channel}`); } async exists(key: string): Promise { diff --git a/src/services/auth.test.ts b/src/services/auth.test.ts index f8e167f..e79dd10 100644 --- a/src/services/auth.test.ts +++ b/src/services/auth.test.ts @@ -153,10 +153,16 @@ describe("auth utils", () => { const first = await exchangeToken(authCode); expect(first).toBeDefined(); + // Mock console.error to suppress expected error message + const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {}); + // Second exchange throws await expect(exchangeToken(authCode)).rejects.toThrow( "Duplicate use of authorization code detected" ); + + // Restore console.error + consoleErrorSpy.mockRestore(); }); it("returns undefined for non-existent code", async () => { diff --git a/src/services/auth.ts b/src/services/auth.ts index f9197c0..425de79 100644 --- a/src/services/auth.ts +++ b/src/services/auth.ts @@ -243,7 +243,7 @@ export async function exchangeToken( const tokenExchange: TokenExchange = JSON.parse(decoded); if (tokenExchange.alreadyUsed) { - console.error("1Duplicate use of authorization code detected; revoking tokens"); + console.error("Duplicate use of authorization code detected; revoking tokens"); await revokeMcpInstallation(tokenExchange.mcpAccessToken); throw new Error("Duplicate use of authorization code detected; tokens revoked"); } diff --git a/src/services/mcp.ts b/src/services/mcp.ts index 9551e6b..04a5e65 100644 --- a/src/services/mcp.ts +++ b/src/services/mcp.ts @@ -95,7 +95,12 @@ enum PromptName { RESOURCE = "resource_prompt", } -export const createMcpServer = () => { +interface McpServerWrapper { + server: Server; + cleanup: () => void; +} + +export const createMcpServer = (): McpServerWrapper => { const server = new Server( { name: "example-servers/everything", diff --git a/src/services/redisTransport.integration.test.ts b/src/services/redisTransport.integration.test.ts new file mode 100644 index 0000000..0acf893 --- /dev/null +++ b/src/services/redisTransport.integration.test.ts @@ -0,0 +1,256 @@ +import { jest } from '@jest/globals'; +import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js'; +import { MockRedisClient, setRedisClient } from '../redis.js'; +import { + startServerListeningToRedis, + redisRelayToMcpServer, + shutdownSession +} from './redisTransport.js'; +import { createMcpServer } from './mcp.js'; +import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; + +describe('Redis Transport Integration', () => { + let mockRedis: MockRedisClient; + + beforeEach(() => { + mockRedis = new MockRedisClient(); + setRedisClient(mockRedis); + jest.resetAllMocks(); + }); + + afterEach(() => { + mockRedis.clear(); + }); + + describe('MCP Initialization Flow', () => { + const sessionId = 'test-init-session'; + + it('should relay initialization request from client to server through Redis', async () => { + // 1. Start the server listening to Redis + const serverWrapper = createMcpServer(); + const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId); + + // 2. Create a mock client transport (simulating the streamable HTTP client side) + const mockClientTransport: Transport = { + onmessage: undefined, + onclose: undefined, + onerror: undefined, + send: jest.fn(() => Promise.resolve()), + close: jest.fn(() => Promise.resolve()), + start: jest.fn(() => Promise.resolve()) + }; + + // 3. Set up the Redis relay (this is what happens in the HTTP handler) + const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport); + + // Track messages received by server + const serverReceivedMessages: JSONRPCMessage[] = []; + const originalServerOnMessage = serverTransport.onmessage; + serverTransport.onmessage = (message, extra) => { + serverReceivedMessages.push(message); + originalServerOnMessage?.(message, extra); + }; + + // 4. Simulate client sending initialization request + const initMessage: JSONRPCMessage = { + jsonrpc: '2.0', + id: 'init-1', + method: 'initialize', + params: { + protocolVersion: '2024-11-05', + capabilities: {}, + clientInfo: { + name: 'test-client', + version: '1.0.0' + } + } + }; + + // Trigger the client transport onmessage (simulates HTTP request) + mockClientTransport.onmessage?.(initMessage); + + // Wait for message to be relayed through Redis + await new Promise(resolve => setTimeout(resolve, 50)); + + // 5. Verify server received the init message + expect(serverReceivedMessages).toHaveLength(1); + expect(serverReceivedMessages[0]).toMatchObject({ + jsonrpc: '2.0', + id: 'init-1', + method: 'initialize' + }); + + // 6. Simulate server responding (this should get relayed back to client) + const initResponse: JSONRPCMessage = { + jsonrpc: '2.0', + id: 'init-1', + result: { + protocolVersion: '2024-11-05', + capabilities: { + tools: {}, + prompts: {}, + resources: {} + }, + serverInfo: { + name: 'example-server', + version: '1.0.0' + } + } + }; + + await serverTransport.send(initResponse, { relatedRequestId: 'init-1' }); + + // Wait for response to be relayed back + await new Promise(resolve => setTimeout(resolve, 50)); + + // 7. Verify client transport received the response + expect(mockClientTransport.send).toHaveBeenCalledWith( + expect.objectContaining({ + jsonrpc: '2.0', + id: 'init-1', + result: expect.objectContaining({ + protocolVersion: '2024-11-05', + serverInfo: expect.objectContaining({ + name: 'example-server' + }) + }) + }), + { relatedRequestId: 'init-1' } + ); + + // Cleanup + await cleanup(); + await shutdownSession(sessionId); + await new Promise(resolve => setTimeout(resolve, 10)); + }); + + it('should handle tools/list request through Redis relay', async () => { + // Set up server and mock client + const serverWrapper = createMcpServer(); + const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId); + + const mockClientTransport: Transport = { + onmessage: undefined, + onclose: undefined, + onerror: undefined, + send: jest.fn(() => Promise.resolve()), + close: jest.fn(() => Promise.resolve()), + start: jest.fn(() => Promise.resolve()) + }; + + const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport); + + // Send tools/list request + const toolsListMessage: JSONRPCMessage = { + jsonrpc: '2.0', + id: 'tools-1', + method: 'tools/list', + params: {} + }; + + mockClientTransport.onmessage?.(toolsListMessage); + + // Wait for processing and response + await new Promise(resolve => setTimeout(resolve, 100)); + + // Verify client received a response with tools + expect(mockClientTransport.send).toHaveBeenCalledWith( + expect.objectContaining({ + jsonrpc: '2.0', + id: 'tools-1', + result: expect.objectContaining({ + tools: expect.any(Array) + }) + }), + undefined + ); + + // Cleanup + await cleanup(); + await shutdownSession(sessionId); + await new Promise(resolve => setTimeout(resolve, 10)); + }); + + it('should handle notifications through Redis relay', async () => { + // Set up server and mock client + const serverWrapper = createMcpServer(); + const serverTransport = await startServerListeningToRedis(serverWrapper, sessionId); + + const mockClientTransport: Transport = { + onmessage: undefined, + onclose: undefined, + onerror: undefined, + send: jest.fn(() => Promise.resolve()), + close: jest.fn(() => Promise.resolve()), + start: jest.fn(() => Promise.resolve()) + }; + + const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport); + + // Send a notification from server (notifications don't have an id) + const notification: JSONRPCMessage = { + jsonrpc: '2.0', + method: 'notifications/message', + params: { + level: 'info', + logger: 'test', + data: 'Test notification' + } + }; + + await serverTransport.send(notification); + + // Wait for notification to be delivered + await new Promise(resolve => setTimeout(resolve, 50)); + + // Verify client received the notification + expect(mockClientTransport.send).toHaveBeenCalledWith( + expect.objectContaining({ + jsonrpc: '2.0', + method: 'notifications/message', + params: expect.objectContaining({ + level: 'info', + data: 'Test notification' + }) + }), + undefined + ); + + // Cleanup + await cleanup(); + await shutdownSession(sessionId); + await new Promise(resolve => setTimeout(resolve, 10)); + }); + + it('should not create response subscriptions for notifications', async () => { + const mockClientTransport: Transport = { + onmessage: undefined, + onclose: undefined, + onerror: undefined, + send: jest.fn(() => Promise.resolve()), + close: jest.fn(() => Promise.resolve()), + start: jest.fn(() => Promise.resolve()) + }; + + const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport); + + // Send a notification (no id field) + const notification: JSONRPCMessage = { + jsonrpc: '2.0', + method: 'notifications/initialized', + params: {} + }; + + mockClientTransport.onmessage?.(notification); + + // Wait for processing + await new Promise(resolve => setTimeout(resolve, 50)); + + // Should not create any response channel subscriptions for notifications + // (we can't easily test this directly, but we can ensure no errors occur) + expect(mockClientTransport.send).not.toHaveBeenCalled(); + + await cleanup(); + }); + }); +}); \ No newline at end of file diff --git a/src/services/redisTransport.test.ts b/src/services/redisTransport.test.ts index 2feb142..5c60e14 100644 --- a/src/services/redisTransport.test.ts +++ b/src/services/redisTransport.test.ts @@ -5,7 +5,8 @@ import { ServerRedisTransport, redisRelayToMcpServer, isLive, - startServerListeningToRedis + startServerListeningToRedis, + shutdownSession } from './redisTransport.js'; import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; import { createMcpServer } from './mcp.js'; @@ -60,6 +61,7 @@ describe('Redis Transport', () => { expect(mockSubscriber).toHaveBeenCalledWith( JSON.stringify({ + type: 'mcp', message: responseMessage, options: { relatedRequestId: 123 } }) @@ -84,6 +86,7 @@ describe('Redis Transport', () => { expect(mockSubscriber).toHaveBeenCalledWith( JSON.stringify({ + type: 'mcp', message: notificationMessage, options: undefined }) @@ -98,6 +101,55 @@ describe('Redis Transport', () => { expect(onCloseMock).toHaveBeenCalled(); }); + + it('should respond to shutdown control messages', async () => { + await transport.start(); + + const onCloseMock = jest.fn(); + transport.onclose = onCloseMock; + + // Send a shutdown control message + await shutdownSession(sessionId); + + // Wait for async processing + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(onCloseMock).toHaveBeenCalled(); + }); + + it('should receive MCP messages from clients and call onmessage', async () => { + const onMessageMock = jest.fn(); + transport.onmessage = onMessageMock; + + await transport.start(); + + // Simulate client sending a message to server + const clientMessage: JSONRPCMessage = { + jsonrpc: '2.0', + id: 'test-req', + method: 'tools/list', + params: {} + }; + + await mockRedis.publish( + `mcp:shttp:toserver:${sessionId}`, + JSON.stringify({ + type: 'mcp', + message: clientMessage, + extra: { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } } + }) + ); + + // Wait for async processing + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(onMessageMock).toHaveBeenCalledWith( + clientMessage, + { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } } + ); + + await transport.close(); + }); }); @@ -110,9 +162,9 @@ describe('Redis Transport', () => { onmessage: undefined, onclose: undefined, onerror: undefined, - send: jest.fn(), - close: jest.fn(), - start: jest.fn() + send: jest.fn(() => Promise.resolve()), + close: jest.fn(() => Promise.resolve()), + start: jest.fn(() => Promise.resolve()) }; }); @@ -128,7 +180,7 @@ describe('Redis Transport', () => { }; // Trigger the onmessage handler - mockTransport.onmessage?.(requestMessage, { authInfo: { userId: 'test' } }); + mockTransport.onmessage?.(requestMessage, { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } }); // Wait a bit for async processing await new Promise(resolve => setTimeout(resolve, 10)); @@ -158,7 +210,7 @@ describe('Redis Transport', () => { }; // Trigger the onmessage handler - mockTransport.onmessage?.(requestMessage, { authInfo: { userId: 'test' } }); + mockTransport.onmessage?.(requestMessage, { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } }); // Wait for subscription setup await new Promise(resolve => setTimeout(resolve, 10)); @@ -173,6 +225,7 @@ describe('Redis Transport', () => { await mockRedis.publish( `mcp:shttp:toclient:${sessionId}:req-456`, JSON.stringify({ + type: 'mcp', message: responseMessage, options: undefined }) @@ -229,24 +282,32 @@ describe('Redis Transport', () => { const sessionId = 'test-session-server'; it('should connect server with ServerRedisTransport', async () => { - const { server } = createMcpServer(); - const connectSpy = jest.spyOn(server, 'connect'); + const serverWrapper = createMcpServer(); + const connectSpy = jest.spyOn(serverWrapper.server, 'connect'); - await startServerListeningToRedis(server, sessionId); + await startServerListeningToRedis(serverWrapper, sessionId); expect(connectSpy).toHaveBeenCalledWith( expect.any(ServerRedisTransport) ); + + // Clean up properly + await shutdownSession(sessionId); + await new Promise(resolve => setTimeout(resolve, 10)); }); - it('should create transport that can send responses', async () => { - const { server } = createMcpServer(); + it('should create transport that can send responses and shutdown properly', async () => { + const serverWrapper = createMcpServer(); - await startServerListeningToRedis(server, sessionId); + const transport = await startServerListeningToRedis(serverWrapper, sessionId); // The server should now be connected and able to handle requests via Redis - // This is tested implicitly by the connection succeeding - expect(server).toBeDefined(); + expect(serverWrapper.server).toBeDefined(); + expect(transport).toBeInstanceOf(ServerRedisTransport); + + // Test proper shutdown + await shutdownSession(sessionId); + await new Promise(resolve => setTimeout(resolve, 10)); }); }); @@ -259,9 +320,9 @@ describe('Redis Transport', () => { onmessage: undefined, onclose: undefined, onerror: undefined, - send: jest.fn(), - close: jest.fn(), - start: jest.fn() + send: jest.fn(() => Promise.resolve()), + close: jest.fn(() => Promise.resolve()), + start: jest.fn(() => Promise.resolve()) }; const cleanup = redisRelayToMcpServer(sessionId, clientTransport); @@ -291,6 +352,7 @@ describe('Redis Transport', () => { // Verify the message was published to server channel expect(serverSubscriber).toHaveBeenCalledWith( JSON.stringify({ + type: 'mcp', message: listToolsRequest, extra: undefined, options: undefined @@ -307,6 +369,7 @@ describe('Redis Transport', () => { await mockRedis.publish( `mcp:shttp:toclient:${sessionId}:integration-req-1`, JSON.stringify({ + type: 'mcp', message: serverResponse, options: undefined }) @@ -321,4 +384,42 @@ describe('Redis Transport', () => { await cleanup(); }); }); + + describe('Control Messages', () => { + const sessionId = 'test-control-session'; + + it('should send shutdown control messages', async () => { + const controlSubscriber = jest.fn(); + await mockRedis.createSubscription( + `mcp:control:${sessionId}`, + controlSubscriber, + jest.fn() + ); + + await shutdownSession(sessionId); + + const callArgs = controlSubscriber.mock.calls[0][0] as string; + const message = JSON.parse(callArgs); + + expect(message.type).toBe('control'); + expect(message.action).toBe('SHUTDOWN'); + expect(typeof message.timestamp).toBe('number'); + }); + + it('should properly shutdown server transport via control message', async () => { + const transport = new ServerRedisTransport(sessionId); + const onCloseMock = jest.fn(); + transport.onclose = onCloseMock; + + await transport.start(); + + // Send shutdown signal + await shutdownSession(sessionId); + + // Wait for async processing + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(onCloseMock).toHaveBeenCalled(); + }); + }); }); \ No newline at end of file diff --git a/src/services/redisTransport.ts b/src/services/redisTransport.ts index 8588190..647fa15 100644 --- a/src/services/redisTransport.ts +++ b/src/services/redisTransport.ts @@ -8,16 +8,25 @@ import { Server } from "@modelcontextprotocol/sdk/server/index.js"; let redisTransportCounter = 0; const notificationStreamId = "__GET_stream"; -interface RedisMessage { - message: JSONRPCMessage; - extra?: MessageExtraInfo; - options?: TransportSendOptions; -} +// Message types for Redis transport +type RedisMessage = + | { + type: 'mcp'; + message: JSONRPCMessage; + extra?: MessageExtraInfo; + options?: TransportSendOptions; + } + | { + type: 'control'; + action: 'SHUTDOWN' | 'PING' | 'STATUS'; + timestamp?: number; + }; function sendToMcpServer(sessionId: string, message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }, options?: TransportSendOptions): Promise { const toServerChannel = getToServerChannel(sessionId); - const redisMessage: RedisMessage = { message, extra, options }; - console.log(`[sendToServerChannel] Publishing to ${toServerChannel}:`, JSON.stringify(redisMessage).substring(0, 100)); + const redisMessage: RedisMessage = { type: 'mcp', message, extra, options }; + console.log(`[sendToMcpServer] Publishing to Redis channel: ${toServerChannel}`); + console.log(`[sendToMcpServer] Message:`, JSON.stringify(message)); return redisClient.publish(toServerChannel, JSON.stringify(redisMessage)); } @@ -29,48 +38,79 @@ function getToClientChannel(sessionId: string, relatedRequestId: string): string return `mcp:shttp:toclient:${sessionId}:${relatedRequestId}`; } +function getControlChannel(sessionId: string): string { + return `mcp:control:${sessionId}`; +} + +function sendControlMessage(sessionId: string, action: 'SHUTDOWN' | 'PING' | 'STATUS'): Promise { + const controlChannel = getControlChannel(sessionId); + const redisMessage: RedisMessage = { + type: 'control', + action, + timestamp: Date.now() + }; + return redisClient.publish(controlChannel, JSON.stringify(redisMessage)); +} + +export async function shutdownSession(sessionId: string): Promise { + return sendControlMessage(sessionId, 'SHUTDOWN'); +} + export async function isLive(sessionId: string): Promise { // Check if the session is live by checking if the key exists in Redis const numSubs = await redisClient.numsub(getToServerChannel(sessionId)); - console.log(`[isLive] Session ${sessionId}: Redis subscribers on ${getToServerChannel(sessionId)} = ${numSubs}`); return numSubs > 0; } export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise { + console.log(`[redisRelayToMcpServer] Setting up relay for session: ${sessionId}`); let redisCleanup: (() => Promise) | undefined = undefined; const cleanup = async () => { + console.log(`[redisRelayToMcpServer] Cleaning up relay for session: ${sessionId}`); // TODO: solve race conditions where we call cleanup while the subscription is being created / before it is created if (redisCleanup) { await redisCleanup(); } } - new Promise((resolve) => { + const messagePromise = new Promise((resolve) => { + console.log(`[redisRelayToMcpServer] Setting up transport.onmessage handler for session: ${sessionId}`); transport.onmessage = async (message, extra) => { + console.log(`[redisRelayToMcpServer] Received message from transport for session ${sessionId}:`, JSON.stringify(message)); + + // First, set up response subscription if needed + if ("id" in message) { + console.log(`[redisRelayToMcpServer] Setting up response subscription for message ID: ${message.id}`); + const toClientChannel = getToClientChannel(sessionId, message.id.toString()); + console.log(`[redisRelayToMcpServer] Subscribing to response channel: ${toClientChannel}`); + + redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => { + console.log(`[redisRelayToMcpServer] Received response from Redis for session ${sessionId}:`, redisMessageJson.substring(0, 200)); + const redisMessage = JSON.parse(redisMessageJson) as RedisMessage; + if (redisMessage.type === 'mcp') { + console.log(`[redisRelayToMcpServer] Sending response back to transport for session ${sessionId}`); + await transport.send(redisMessage.message, redisMessage.options); + console.log(`[redisRelayToMcpServer] Response sent to transport for session ${sessionId}`); + } + }, (error) => { + console.error(`[redisRelayToMcpServer] Error in response subscription for session ${sessionId}:`, error); + transport.onerror?.(error); + }); + console.log(`[redisRelayToMcpServer] Response subscription established for session ${sessionId}`); + } else { + console.log(`[redisRelayToMcpServer] Message is notification (no ID) for session ${sessionId}, skipping response subscription`); + } + + // Now send the message to the MCP server await sendToMcpServer(sessionId, message, extra); + console.log(`[redisRelayToMcpServer] Sent message to MCP server for session ${sessionId}`); resolve(message); } - }).then(async (message) => { - // check for request id in the message - if (!("id" in message)) { - // if no id, it's a notification, so we return - return cleanup; - } - // otherwise we subscribe to the response channel - const toClientChannel = getToClientChannel(sessionId, message.id.toString()); + }); - console.log(`[redisRelayToMcpServer] Subscribing to ${toClientChannel} for response to request ${message.id}`); - - redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => { - const redisMessage = JSON.parse(redisMessageJson) as RedisMessage; - await transport.send(redisMessage.message, redisMessage.options); - }, (error) => { - console.error(`[redisRelayToMcpServer] Error in Redis subscriber for ${toClientChannel}:`, error); - transport.onerror?.(error); - }); - }).catch((error) => { - console.error(`[redisRelayToMcpServer] Error setting up Redis relay for session ${sessionId}:`, error); + messagePromise.catch((error) => { + console.error(`[redisRelayToMcpServer] Error setting up relay for session ${sessionId}:`, error); transport.onerror?.(error); cleanup(); }); @@ -83,6 +123,9 @@ export function redisRelayToMcpServer(sessionId: string, transport: Transport): export class ServerRedisTransport implements Transport { private counter: number; private _sessionId: string; + private controlCleanup?: (() => Promise); + private serverCleanup?: (() => Promise); + private shouldShutdown = false; onclose?: (() => void) | undefined; onerror?: ((error: Error) => void) | undefined; @@ -94,45 +137,108 @@ export class ServerRedisTransport implements Transport { } async start(): Promise { - console.log(`[ServerRedisTransport.${this.counter}.start] Starting server transport for session: ${this._sessionId}`); - // Server transport doesn't need to subscribe to anything - it only sends responses + + // Subscribe to MCP messages from clients + const serverChannel = getToServerChannel(this._sessionId); + console.log(`[ServerRedisTransport.${this.counter}] Subscribing to server channel: ${serverChannel}`); + this.serverCleanup = await redisClient.createSubscription( + serverChannel, + (messageJson) => { + console.log(`[ServerRedisTransport.${this.counter}] Received message from Redis:`, messageJson.substring(0, 200)); + const redisMessage = JSON.parse(messageJson) as RedisMessage; + if (redisMessage.type === 'mcp') { + console.log(`[ServerRedisTransport.${this.counter}] Processing MCP message:`, JSON.stringify(redisMessage.message)); + this.onmessage?.(redisMessage.message, redisMessage.extra); + console.log(`[ServerRedisTransport.${this.counter}] MCP message processed`); + } + }, + (error) => { + console.error(`[ServerRedisTransport.${this.counter}] Server channel error:`, error); + this.onerror?.(error); + } + ); + + // Subscribe to control messages for shutdown + const controlChannel = getControlChannel(this._sessionId); + this.controlCleanup = await redisClient.createSubscription( + controlChannel, + (messageJson) => { + const redisMessage = JSON.parse(messageJson) as RedisMessage; + if (redisMessage.type === 'control') { + if (redisMessage.action === 'SHUTDOWN') { + this.shouldShutdown = true; + this.close(); + } + } + }, + (error) => { + this.onerror?.(error); + } + ); + } async send(message: JSONRPCMessage, options?: TransportSendOptions): Promise { - const relatedRequestId = options?.relatedRequestId?.toString() ?? notificationStreamId; + const relatedRequestId = options?.relatedRequestId?.toString() ?? "id" in message ? message.id : notificationStreamId; const channel = getToClientChannel(this._sessionId, relatedRequestId) + console.log(`[ServerRedisTransport.${this.counter}] Sending message to channel: ${channel}`); + console.log(`[ServerRedisTransport.${this.counter}] Message:`, JSON.stringify(message)); + console.log(`[ServerRedisTransport.${this.counter}] Options:`, JSON.stringify(options)); - const redisMessage: RedisMessage = { message, options }; + const redisMessage: RedisMessage = { type: 'mcp', message, options }; const messageStr = JSON.stringify(redisMessage); - console.log(`[ServerRedisTransport.${this.counter}.send] Publishing to ${channel}:`, messageStr.substring(0, 100)); await redisClient.publish(channel, messageStr); - console.log(`[ServerRedisTransport.${this.counter}.send] Published successfully to ${channel}`); + console.log(`[ServerRedisTransport.${this.counter}] Message published to Redis`); } async close(): Promise { - console.log(`[ServerRedisTransport.${this.counter}.close] Closing server transport for session: ${this._sessionId}`); + + // Clean up server message subscription + if (this.serverCleanup) { + await this.serverCleanup(); + this.serverCleanup = undefined; + } + + // Clean up control message subscription + if (this.controlCleanup) { + await this.controlCleanup(); + this.controlCleanup = undefined; + } + this.onclose?.(); } } -export async function startServerListeningToRedis(server: Server, sessionId: string) { - console.log(`[startServerListeningToRedis] Starting background server for session ${sessionId}`); +export async function startServerListeningToRedis(serverWrapper: { server: Server; cleanup: () => void }, sessionId: string): Promise { const serverRedisTransport = new ServerRedisTransport(sessionId); + // Add cleanup callback to the transport + const originalClose = serverRedisTransport.close.bind(serverRedisTransport); + serverRedisTransport.close = async () => { + serverWrapper.cleanup(); + await originalClose(); + }; + // The server.connect() will call start() on the transport - await server.connect(serverRedisTransport) - console.log(`[startServerListeningToRedis] Background server connected for session ${sessionId}`); + await serverWrapper.server.connect(serverRedisTransport) + + return serverRedisTransport; } export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, cleanup: () => Promise}> { + console.log(`[getFirstShttpTransport] Creating transport for session: ${sessionId}`); const shttpTransport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => sessionId, enableJsonResponse: true, // Enable JSON response mode }); + console.log(`[getFirstShttpTransport] Created transport, sessionId: ${shttpTransport.sessionId}`); + // Use the new request-id based relay approach + console.log(`[getFirstShttpTransport] Setting up Redis relay for session: ${sessionId}`); const cleanup = redisRelayToMcpServer(sessionId, shttpTransport); + console.log(`[getFirstShttpTransport] Transport setup complete for session: ${sessionId}`); return { shttpTransport, cleanup }; }