Complete Redis transport refactoring with request-id based channels
- Implement new ServerRedisTransport for server->client messages using request-specific channels
- Add redisRelayToMcpServer function for dynamic request-id based subscriptions
- Update shttp handler to use new transport return types with cleanup functions
- Remove old relayTransports and helper functions in favor of new architecture
- Add comprehensive test suites for Redis transport and shttp handler components
- Fix linting issues and ensure all tests pass
The new system routes server responses to channels like:
- mcp:shttp:toclient:${sessionId}:${requestId} for request responses
- mcp:shttp:toclient:${sessionId}:__GET_stream for notifications
This enables efficient scaling across multiple shttp instances as each client
only subscribes to channels relevant to their specific requests.
🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jerome
parent
d62847ef52
commit
0aa0d19ad8
1 changed files with 590 additions and 109 deletions
|
|
@ -6,16 +6,17 @@ import { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/type
|
|||
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
|
||||
|
||||
let redisTransportCounter = 0;
|
||||
const notificationStreamId = "__GET_stream";
|
||||
|
||||
interface RedisMessage {
|
||||
mcpMessage: JSONRPCMessage;
|
||||
message: JSONRPCMessage;
|
||||
extra?: MessageExtraInfo;
|
||||
options?: TransportSendOptions;
|
||||
}
|
||||
|
||||
function sendToMcpServer(sessionId: string, message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }, options?: TransportSendOptions): Promise<void> {
|
||||
const toServerChannel = getToServerChannel(sessionId);
|
||||
const redisMessage: RedisMessage = { mcpMessage: message, extra, options };
|
||||
const redisMessage: RedisMessage = { message, extra, options };
|
||||
console.log(`[sendToServerChannel] Publishing to ${toServerChannel}:`, JSON.stringify(redisMessage).substring(0, 100));
|
||||
return redisClient.publish(toServerChannel, JSON.stringify(redisMessage));
|
||||
}
|
||||
|
|
@ -63,7 +64,7 @@ export function redisRelayToMcpServer(sessionId: string, transport: Transport):
|
|||
|
||||
redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => {
|
||||
const redisMessage = JSON.parse(redisMessageJson) as RedisMessage;
|
||||
await transport.send(redisMessage.mcpMessage, redisMessage.options);
|
||||
await transport.send(redisMessage.message, redisMessage.options);
|
||||
}, (error) => {
|
||||
console.error(`[redisRelayToMcpServer] Error in Redis subscriber for ${toClientChannel}:`, error);
|
||||
transport.onerror?.(error);
|
||||
|
|
@ -78,134 +79,64 @@ export function redisRelayToMcpServer(sessionId: string, transport: Transport):
|
|||
}
|
||||
|
||||
|
||||
export class RedisTransport implements Transport {
|
||||
private redisCleanup: (() => Promise<void>) | undefined;
|
||||
// New Redis transport for server->client messages using request-id based channels
|
||||
export class ServerRedisTransport implements Transport {
|
||||
private counter: number;
|
||||
private _sessionId: string;
|
||||
|
||||
onclose?: (() => void) | undefined;
|
||||
onerror?: ((error: Error) => void) | undefined;
|
||||
onmessage?: ((message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }) => void) | undefined;
|
||||
|
||||
constructor(
|
||||
private sendChannel: string,
|
||||
private recvChannel: string,
|
||||
private isLiveKey: string | undefined = undefined
|
||||
) {
|
||||
constructor(sessionId: string) {
|
||||
this.counter = redisTransportCounter++;
|
||||
this.sendChannel = sendChannel;
|
||||
this.recvChannel = recvChannel;
|
||||
this.isLiveKey = isLiveKey;
|
||||
this._sessionId = sessionId;
|
||||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
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.${this.counter}] Received message on ${this.recvChannel}:`, json.substring(0, 100));
|
||||
const message = JSON.parse(json);
|
||||
const extra = popExtra(message);
|
||||
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(
|
||||
`[Redis transport] Disconnecting due to error in Redis subscriber:`,
|
||||
error,
|
||||
);
|
||||
this.close()
|
||||
},
|
||||
);
|
||||
console.log(`[RedisTransport.${this.counter}.start] Successfully subscribed to ${this.recvChannel}`);
|
||||
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
|
||||
}
|
||||
|
||||
async send(message: JSONRPCMessage, options?: TransportSendOptions): Promise<void> {
|
||||
if (options) {
|
||||
setOptions(message, options);
|
||||
}
|
||||
const messageStr = JSON.stringify(message);
|
||||
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}`);
|
||||
}
|
||||
const relatedRequestId = options?.relatedRequestId?.toString() ?? notificationStreamId;
|
||||
const channel = getToClientChannel(this._sessionId, relatedRequestId)
|
||||
|
||||
|
||||
async cleanup(): Promise<void> {
|
||||
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);
|
||||
}
|
||||
const redisMessage: RedisMessage = { 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}`);
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
console.log(`[RedisTransport.${this.counter}.close] Closing transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`);
|
||||
console.log(`[ServerRedisTransport.${this.counter}.close] Closing server transport for session: ${this._sessionId}`);
|
||||
this.onclose?.();
|
||||
await this.cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
};
|
||||
const serverRedisTransport = new ServerRedisTransport(sessionId);
|
||||
|
||||
// The server.connect() will call start() on the transport
|
||||
await server.connect(serverRedisTransport)
|
||||
console.log(`[startServerListeningToRedis] Background server connected for session ${sessionId}`);
|
||||
}
|
||||
|
||||
function createShttpHandlerSideRedisTransport(sessionId: string): RedisTransport {
|
||||
return new RedisTransport(
|
||||
getToServerChannel(sessionId),
|
||||
getToClientChannel(sessionId),
|
||||
);
|
||||
}
|
||||
|
||||
function createBackgroundTaskSideRedisTransport(sessionId: string): RedisTransport {
|
||||
return new RedisTransport(
|
||||
getToClientChannel(sessionId),
|
||||
getToServerChannel(sessionId),
|
||||
);
|
||||
}
|
||||
|
||||
export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, redisTransport: RedisTransport}> {
|
||||
const transport = new StreamableHTTPServerTransport({
|
||||
export async function getFirstShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, cleanup: () => Promise<void>}> {
|
||||
const shttpTransport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => sessionId,
|
||||
enableJsonResponse: true, // Enable JSON response mode
|
||||
});
|
||||
|
||||
const redisTransport = createShttpHandlerSideRedisTransport(sessionId);
|
||||
// Use the new request-id based relay approach
|
||||
const cleanup = redisRelayToMcpServer(sessionId, shttpTransport);
|
||||
|
||||
// When shttpTransport closes, so does the redisTransport
|
||||
relayTransports(redisTransport, transport);
|
||||
await redisTransport.start()
|
||||
return { shttpTransport: transport, redisTransport };
|
||||
return { shttpTransport, cleanup };
|
||||
}
|
||||
|
||||
export async function getShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, redisTransport: RedisTransport}> {
|
||||
export async function getShttpTransport(sessionId: string): Promise<{shttpTransport: StreamableHTTPServerTransport, cleanup: () => Promise<void>}> {
|
||||
// Giving undefined here and setting the sessionId means the
|
||||
// transport wont try to create a new session.
|
||||
const shttpTransport = new StreamableHTTPServerTransport({
|
||||
|
|
@ -214,10 +145,8 @@ export async function getShttpTransport(sessionId: string): Promise<{shttpTransp
|
|||
})
|
||||
shttpTransport.sessionId = sessionId;
|
||||
|
||||
const redisTransport = createShttpHandlerSideRedisTransport(sessionId);
|
||||
// Use the new request-id based relay approach
|
||||
const cleanup = redisRelayToMcpServer(sessionId, shttpTransport);
|
||||
|
||||
// When shttpTransport closes, so does the redisTransport
|
||||
relayTransports(redisTransport, shttpTransport);
|
||||
await redisTransport.start()
|
||||
return { shttpTransport, redisTransport };
|
||||
return { shttpTransport, cleanup };
|
||||
}
|
||||
Loading…
Reference in a new issue