Fix Redis transport refactoring with control message system
- 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 <noreply@anthropic.com>
This commit is contained in:
Jerome
parent
0aa0d19ad8
commit
de8c4a4ed2
8 changed files with 552 additions and 62 deletions
|
|
@ -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));
|
||||
|
|
|
|||
20
src/redis.ts
20
src/redis.ts
|
|
@ -104,18 +104,34 @@ export class RedisClientImpl implements RedisClient {
|
|||
onMessage: (message: string) => void,
|
||||
onError: (error: Error) => void,
|
||||
): Promise<() => Promise<void>> {
|
||||
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<void> {
|
||||
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<boolean> {
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
256
src/services/redisTransport.integration.test.ts
Normal file
256
src/services/redisTransport.integration.test.ts
Normal file
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
return sendControlMessage(sessionId, 'SHUTDOWN');
|
||||
}
|
||||
|
||||
export async function isLive(sessionId: string): Promise<boolean> {
|
||||
// Check if the session is live by checking if the key exists in Redis
|
||||
const numSubs = await redisClient.numsub(getToServerChannel(sessionId));
|
||||
console.log(`[isLive] Session ${sessionId}: Redis subscribers on ${getToServerChannel(sessionId)} = ${numSubs}`);
|
||||
return numSubs > 0;
|
||||
}
|
||||
|
||||
|
||||
export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise<void> {
|
||||
console.log(`[redisRelayToMcpServer] Setting up relay for session: ${sessionId}`);
|
||||
let redisCleanup: (() => Promise<void>) | 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<JSONRPCMessage>((resolve) => {
|
||||
const messagePromise = new Promise<JSONRPCMessage>((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<void>);
|
||||
private serverCleanup?: (() => Promise<void>);
|
||||
private shouldShutdown = false;
|
||||
|
||||
onclose?: (() => void) | undefined;
|
||||
onerror?: ((error: Error) => void) | undefined;
|
||||
|
|
@ -94,45 +137,108 @@ export class ServerRedisTransport implements Transport {
|
|||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
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<ServerRedisTransport> {
|
||||
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<void>}> {
|
||||
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 };
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue