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;
|
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)
|
||||||
|
|
|
||||||
|
|
@ -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;
|
||||||
}
|
}
|
||||||
Loading…
Reference in a new issue