Wired up notifications
This commit is contained in:
parent
29f5bde8b2
commit
728c3f66ea
2 changed files with 37 additions and 29 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
Loading…
Reference in a new issue