Part way through refactoring the redis transport
This commit is contained in:
parent
c2119ed9f9
commit
d62847ef52
1 changed files with 58 additions and 116 deletions
|
|
@ -2,17 +2,30 @@ import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/
|
|||
import { redisClient } from "../redis.js";
|
||||
import { Transport, TransportSendOptions } from "@modelcontextprotocol/sdk/shared/transport.js";
|
||||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||
import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
|
||||
import { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js";
|
||||
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
|
||||
|
||||
let redisTransportCounter = 0;
|
||||
|
||||
interface RedisMessage {
|
||||
mcpMessage: 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 };
|
||||
console.log(`[sendToServerChannel] Publishing to ${toServerChannel}:`, JSON.stringify(redisMessage).substring(0, 100));
|
||||
return redisClient.publish(toServerChannel, JSON.stringify(redisMessage));
|
||||
}
|
||||
|
||||
function getToServerChannel(sessionId: string): string {
|
||||
return `mcp:shttp:toserver:${sessionId}`;
|
||||
}
|
||||
|
||||
function getToClientChannel(sessionId: string): string {
|
||||
return `mcp:shttp:toclient:${sessionId}`;
|
||||
function getToClientChannel(sessionId: string, relatedRequestId: string): string {
|
||||
return `mcp:shttp:toclient:${sessionId}:${relatedRequestId}`;
|
||||
}
|
||||
|
||||
export async function isLive(sessionId: string): Promise<boolean> {
|
||||
|
|
@ -23,6 +36,48 @@ export async function isLive(sessionId: string): Promise<boolean> {
|
|||
}
|
||||
|
||||
|
||||
export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise<void> {
|
||||
let redisCleanup: (() => Promise<void>) | undefined = undefined;
|
||||
const cleanup = async () => {
|
||||
// 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) => {
|
||||
transport.onmessage = async (message, extra) => {
|
||||
await sendToMcpServer(sessionId, message, extra);
|
||||
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.mcpMessage, 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);
|
||||
transport.onerror?.(error);
|
||||
cleanup();
|
||||
});
|
||||
|
||||
return cleanup;
|
||||
}
|
||||
|
||||
|
||||
export class RedisTransport implements Transport {
|
||||
private redisCleanup: (() => Promise<void>) | undefined;
|
||||
private counter: number;
|
||||
|
|
@ -42,8 +97,6 @@ export class RedisTransport implements Transport {
|
|||
this.isLiveKey = isLiveKey;
|
||||
}
|
||||
|
||||
|
||||
|
||||
async start(): Promise<void> {
|
||||
console.log(`[RedisTransport.${this.counter}.start] Starting transport - send: ${this.sendChannel}, recv: ${this.recvChannel}`);
|
||||
if (this.redisCleanup) {
|
||||
|
|
@ -108,117 +161,6 @@ export class RedisTransport implements Transport {
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
function setExtra(message: JSONRPCMessage, extra: { authInfo?: AuthInfo; } | undefined): void {
|
||||
if (!extra) {
|
||||
return;
|
||||
}
|
||||
if ("result" in message && typeof message.result === 'object') {
|
||||
if (!message.result._meta) {
|
||||
message.result._meta = {};
|
||||
}
|
||||
message.result._meta.extra = extra;
|
||||
}
|
||||
if ("params" in message && typeof message.params === 'object') {
|
||||
if (!message.params._meta) {
|
||||
message.params._meta = {};
|
||||
}
|
||||
message.params._meta.extra = extra;
|
||||
}
|
||||
}
|
||||
|
||||
function setOptions(message: JSONRPCMessage, options: TransportSendOptions): void {
|
||||
if ("result" in message && typeof message.result === 'object') {
|
||||
if (!message.result._meta) {
|
||||
message.result._meta = {};
|
||||
}
|
||||
message.result._meta.options = options;
|
||||
}
|
||||
if ("params" in message && typeof message.params === 'object') {
|
||||
if (!message.params._meta) {
|
||||
message.params._meta = {};
|
||||
}
|
||||
message.params._meta.options = options;
|
||||
}
|
||||
}
|
||||
|
||||
function popExtra(message: JSONRPCMessage): { authInfo?: AuthInfo; } | undefined {
|
||||
if ("params" in message && typeof message.params === 'object' && message.params._meta) {
|
||||
const extra = message.params._meta.extra as { authInfo?: AuthInfo; } | undefined;
|
||||
if (extra) {
|
||||
delete message.params._meta.extra;
|
||||
return extra;
|
||||
}
|
||||
}
|
||||
if ("result" in message && typeof message.result === 'object' && message.result._meta) {
|
||||
const extra = message.result._meta.extra as { authInfo?: AuthInfo; } | undefined;
|
||||
if (extra) {
|
||||
delete message.result._meta.extra;
|
||||
return extra;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function popOptions(message: JSONRPCMessage): TransportSendOptions | undefined {
|
||||
if ("params" in message && typeof message.params === 'object' && message.params._meta) {
|
||||
const options = message.params._meta.options as TransportSendOptions | undefined;
|
||||
if (options) {
|
||||
delete message.params._meta.options;
|
||||
return options;
|
||||
}
|
||||
}
|
||||
if ("result" in message && typeof message.result === 'object' && message.result._meta) {
|
||||
const options = message.result._meta.options as TransportSendOptions | undefined;
|
||||
if (options) {
|
||||
delete message.result._meta.options;
|
||||
return options;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
|
||||
function relayTransports(transport1: Transport, transport2: Transport): void {
|
||||
console.log(`[relayTransports] Setting up relay between transports`);
|
||||
transport1.onmessage = (message, extra) => {
|
||||
console.log(`[relay] transport1 -> transport2:`, JSON.stringify(message).substring(0, 100));
|
||||
setExtra(message, extra);
|
||||
const options = popOptions(message);
|
||||
transport2.send(message, options)
|
||||
};
|
||||
transport2.onmessage = (message, extra) => {
|
||||
console.log(`[relay] transport2 -> transport1:`, JSON.stringify(message).substring(0, 100));
|
||||
setExtra(message, extra);
|
||||
const options = popOptions(message);
|
||||
transport1.send(message, options)
|
||||
};
|
||||
|
||||
transport1.onerror = (error) => {
|
||||
transport2.onerror?.(error);
|
||||
};
|
||||
transport2.onerror = (error) => {
|
||||
transport1.onerror?.(error);
|
||||
};
|
||||
|
||||
// Prevent circular close calls
|
||||
let closing = false;
|
||||
|
||||
transport1.onclose = () => {
|
||||
if (!closing) {
|
||||
closing = true;
|
||||
transport2.close().catch(console.error);
|
||||
}
|
||||
};
|
||||
transport2.onclose = () => {
|
||||
if (!closing) {
|
||||
closing = true;
|
||||
transport1.close().catch(console.error);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
export async function startServerListeningToRedis(server: Server, sessionId: string) {
|
||||
console.log(`[startServerListeningToRedis] Starting background server for session ${sessionId}`);
|
||||
const serverRedisTransport = createBackgroundTaskSideRedisTransport(sessionId)
|
||||
|
|
|
|||
Loading…
Reference in a new issue