Wired up notifications

This commit is contained in:
Jerome 2025-07-09 16:24:15 +01:00
commit 728c3f66ea
No known key found for this signature in database
2 changed files with 37 additions and 29 deletions

View file

@ -42,6 +42,8 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
return; return;
} }
const isGetRequest = req.method === 'GET';
// incorrect session for the authed user, return 401 // incorrect session for the authed user, return 401
if (sessionId) { if (sessionId) {
if (!(await isSessionOwnedBy(sessionId, userId))) { if (!(await isSessionOwnedBy(sessionId, userId))) {
@ -49,7 +51,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
return; return;
} }
// Reuse existing transport for owned session // Reuse existing transport for owned session
shttpTransport = await getShttpTransport(sessionId, onsessionclosed); shttpTransport = await getShttpTransport(sessionId, onsessionclosed, isGetRequest);
} else if (isInitializeRequest(req.body)) { } else if (isInitializeRequest(req.body)) {
// New initialization request - use JSON response mode // New initialization request - use JSON response mode
const onsessioninitialized = async (sessionId: string) => { const onsessioninitialized = async (sessionId: string) => {
@ -69,7 +71,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) {
onsessionclosed, onsessionclosed,
onsessioninitialized, onsessioninitialized,
}); });
shttpTransport.onclose = redisRelayToMcpServer(sessionId, shttpTransport); shttpTransport.onclose = await redisRelayToMcpServer(sessionId, shttpTransport);
} else { } else {
// Invalid request - no session ID and not initialization request // Invalid request - no session ID and not initialization request
res.status(400) res.status(400)

View file

@ -81,7 +81,7 @@ export async function isSessionOwnedBy(sessionId: string, userId: string): Promi
} }
export function redisRelayToMcpServer(sessionId: string, transport: Transport): () => Promise<void> { export async function redisRelayToMcpServer(sessionId: string, transport: Transport, isGetRequest: boolean = false): Promise<() => Promise<void>> {
let redisCleanup: (() => Promise<void>) | undefined = undefined; let redisCleanup: (() => Promise<void>) | undefined = undefined;
const cleanup = async () => { const cleanup = async () => {
// TODO: solve race conditions where we call cleanup while the subscription is being created / before it is created // 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<JSONRPCMessage>((resolve) => { const subscribe = async (requestId: string) => {
transport.onmessage = async (message, extra) => { const toClientChannel = getToClientChannel(sessionId, requestId);
// First, set up response subscription if needed
if ("id" in message) {
const toClientChannel = getToClientChannel(sessionId, message.id.toString());
redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => { redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => {
const redisMessage = JSON.parse(redisMessageJson) as RedisMessage; const redisMessage = JSON.parse(redisMessageJson) as RedisMessage;
if (redisMessage.type === 'mcp') { if (redisMessage.type === 'mcp') {
await transport.send(redisMessage.message, redisMessage.options); await transport.send(redisMessage.message, redisMessage.options);
}
}, (error) => {
transport.onerror?.(error);
});
} }
}, (error) => {
transport.onerror?.(error);
});
}
// Now send the message to the MCP server if (isGetRequest) {
await sendToMcpServer(sessionId, message, extra); await subscribe(notificationStreamId);
resolve(message); } else {
} const messagePromise = new Promise<JSONRPCMessage>((resolve) => {
}); transport.onmessage = async (message, extra) => {
// First, set up response subscription if needed
messagePromise.catch((error) => { if ("id" in message) {
transport.onerror?.(error); await subscribe(message.id.toString());
cleanup(); }
}); // Now send the message to the MCP server
await sendToMcpServer(sessionId, message, extra);
resolve(message);
}
});
messagePromise.catch((error) => {
transport.onerror?.(error);
cleanup();
});
}
return cleanup; return cleanup;
} }
@ -202,7 +208,7 @@ export class ServerRedisTransport implements Transport {
} }
} }
export async function getShttpTransport(sessionId: string, onsessionclosed: (sessionId: string) => void | Promise<void>): Promise<StreamableHTTPServerTransport> { export async function getShttpTransport(sessionId: string, onsessionclosed: (sessionId: string) => void | Promise<void>, isGetRequest: boolean = false): Promise<StreamableHTTPServerTransport> {
// Giving undefined here and setting the sessionId means the // Giving undefined here and setting the sessionId means the
// transport wont try to create a new session. // transport wont try to create a new session.
const shttpTransport = new StreamableHTTPServerTransport({ const shttpTransport = new StreamableHTTPServerTransport({
@ -212,7 +218,7 @@ export async function getShttpTransport(sessionId: string, onsessionclosed: (ses
shttpTransport.sessionId = sessionId; shttpTransport.sessionId = sessionId;
// Use the new request-id based relay approach // Use the new request-id based relay approach
const cleanup = redisRelayToMcpServer(sessionId, shttpTransport); const cleanup = await redisRelayToMcpServer(sessionId, shttpTransport, isGetRequest);
shttpTransport.onclose = cleanup; shttpTransport.onclose = cleanup;
return shttpTransport; return shttpTransport;
} }