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;
}
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)

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;
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<JSONRPCMessage>((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<JSONRPCMessage>((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<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
// 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;
}