diff --git a/src/handlers/shttp.ts b/src/handlers/shttp.ts index 5ee1e10..d52bd18 100644 --- a/src/handlers/shttp.ts +++ b/src/handlers/shttp.ts @@ -42,6 +42,8 @@ export async function handleStreamableHTTP(req: Request, res: Response) { return; } + const isGetRequest = req.method === 'GET'; + // incorrect session for the authed user, return 401 if (sessionId) { if (!(await isSessionOwnedBy(sessionId, userId))) { @@ -49,7 +51,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) { return; } // Reuse existing transport for owned session - shttpTransport = await getShttpTransport(sessionId, onsessionclosed); + shttpTransport = await getShttpTransport(sessionId, onsessionclosed, isGetRequest); } else if (isInitializeRequest(req.body)) { // New initialization request - use JSON response mode const onsessioninitialized = async (sessionId: string) => { @@ -69,7 +71,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) { onsessionclosed, onsessioninitialized, }); - shttpTransport.onclose = redisRelayToMcpServer(sessionId, shttpTransport); + shttpTransport.onclose = await redisRelayToMcpServer(sessionId, shttpTransport); } else { // Invalid request - no session ID and not initialization request res.status(400) diff --git a/src/services/redisTransport.ts b/src/services/redisTransport.ts index 8697f2b..10ee3e0 100644 --- a/src/services/redisTransport.ts +++ b/src/services/redisTransport.ts @@ -81,7 +81,7 @@ export async function isSessionOwnedBy(sessionId: string, userId: string): Promi } -export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise { +export async function redisRelayToMcpServer(sessionId: string, transport: Transport, isGetRequest: boolean = false): Promise<() => Promise> { let redisCleanup: (() => Promise) | undefined = undefined; const cleanup = async () => { // TODO: solve race conditions where we call cleanup while the subscription is being created / before it is created @@ -90,33 +90,39 @@ export function redisRelayToMcpServer(sessionId: string, transport: Transport): } } - const messagePromise = new Promise((resolve) => { - transport.onmessage = async (message, extra) => { - // First, set up response subscription if needed - if ("id" in message) { - const toClientChannel = getToClientChannel(sessionId, message.id.toString()); + const subscribe = async (requestId: string) => { + const toClientChannel = getToClientChannel(sessionId, requestId); - redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => { - const redisMessage = JSON.parse(redisMessageJson) as RedisMessage; - if (redisMessage.type === 'mcp') { - await transport.send(redisMessage.message, redisMessage.options); - } - }, (error) => { - transport.onerror?.(error); - }); + redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => { + const redisMessage = JSON.parse(redisMessageJson) as RedisMessage; + if (redisMessage.type === 'mcp') { + await transport.send(redisMessage.message, redisMessage.options); } - - // Now send the message to the MCP server - await sendToMcpServer(sessionId, message, extra); - resolve(message); - } - }); - - messagePromise.catch((error) => { - transport.onerror?.(error); - cleanup(); - }); + }, (error) => { + transport.onerror?.(error); + }); + } + if (isGetRequest) { + await subscribe(notificationStreamId); + } else { + const messagePromise = new Promise((resolve) => { + transport.onmessage = async (message, extra) => { + // First, set up response subscription if needed + if ("id" in message) { + await subscribe(message.id.toString()); + } + // Now send the message to the MCP server + await sendToMcpServer(sessionId, message, extra); + resolve(message); + } + }); + + messagePromise.catch((error) => { + transport.onerror?.(error); + cleanup(); + }); + } return cleanup; } @@ -202,7 +208,7 @@ export class ServerRedisTransport implements Transport { } } -export async function getShttpTransport(sessionId: string, onsessionclosed: (sessionId: string) => void | Promise): Promise { +export async function getShttpTransport(sessionId: string, onsessionclosed: (sessionId: string) => void | Promise, isGetRequest: boolean = false): Promise { // Giving undefined here and setting the sessionId means the // transport wont try to create a new session. const shttpTransport = new StreamableHTTPServerTransport({ @@ -212,7 +218,7 @@ export async function getShttpTransport(sessionId: string, onsessionclosed: (ses shttpTransport.sessionId = sessionId; // Use the new request-id based relay approach - const cleanup = redisRelayToMcpServer(sessionId, shttpTransport); + const cleanup = await redisRelayToMcpServer(sessionId, shttpTransport, isGetRequest); shttpTransport.onclose = cleanup; return shttpTransport; } \ No newline at end of file