100 lines
3 KiB
TypeScript
100 lines
3 KiB
TypeScript
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
|
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
|
|
import contentType from "content-type";
|
|
import { Request, Response } from "express";
|
|
import { redisClient } from "../redis.js";
|
|
import { createMcpServer } from "../services/mcp.js";
|
|
import { logger, logMcpMessage } from "../utils/logger.js";
|
|
|
|
declare module "express-serve-static-core" {
|
|
interface Request {
|
|
/**
|
|
* Information about the validated access token, if the `requireBearerAuth` middleware was used.
|
|
*/
|
|
auth?: AuthInfo;
|
|
}
|
|
}
|
|
|
|
function redisChannelForSession(sessionId: string): string {
|
|
return `mcp:${sessionId}`;
|
|
}
|
|
|
|
export async function handleSSEConnection(req: Request, res: Response) {
|
|
const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer();
|
|
const transport = new SSEServerTransport("/message", res);
|
|
logger.info('Received MCP SSE connection', {
|
|
sessionId: transport.sessionId
|
|
});
|
|
|
|
const redisCleanup = await redisClient.createSubscription(
|
|
redisChannelForSession(transport.sessionId),
|
|
(json) => {
|
|
// TODO handle DELETE messages
|
|
// TODO set timeout to kill the session
|
|
|
|
const message = JSON.parse(json);
|
|
logMcpMessage(message, transport.sessionId);
|
|
transport.handleMessage(message).catch((error) => {
|
|
logger.error('Error handling message', error as Error, {
|
|
sessionId: transport.sessionId
|
|
});
|
|
});
|
|
},
|
|
(error) => {
|
|
logger.error('Disconnecting due to error in Redis subscriber', error as Error, {
|
|
sessionId: transport.sessionId
|
|
});
|
|
transport
|
|
.close()
|
|
.catch((error) =>
|
|
logger.error('Error closing transport', error as Error, {
|
|
sessionId: transport.sessionId
|
|
}),
|
|
);
|
|
},
|
|
);
|
|
|
|
const cleanup = () => {
|
|
void mcpCleanup();
|
|
redisCleanup().catch((error) =>
|
|
logger.error('Error disconnecting Redis subscriber', error as Error, {
|
|
sessionId: transport.sessionId
|
|
}),
|
|
);
|
|
}
|
|
|
|
// Clean up Redis subscription when the connection closes
|
|
mcpServer.onclose = cleanup
|
|
|
|
logger.info('Listening on Redis channel', {
|
|
sessionId: transport.sessionId,
|
|
channel: redisChannelForSession(transport.sessionId)
|
|
});
|
|
await mcpServer.connect(transport);
|
|
}
|
|
|
|
export async function handleMessage(req: Request, res: Response) {
|
|
const sessionId = req.query.sessionId;
|
|
let body: string;
|
|
try {
|
|
if (typeof sessionId !== "string") {
|
|
throw new Error("Only one sessionId allowed");
|
|
}
|
|
|
|
const ct = contentType.parse(req.headers["content-type"] ?? "");
|
|
if (ct.type !== "application/json") {
|
|
throw new Error(`Unsupported content-type: ${ct}`);
|
|
}
|
|
|
|
body = JSON.stringify(req.body);
|
|
} catch (error) {
|
|
res.status(400).json(error);
|
|
logger.error('Bad POST request', error as Error, {
|
|
sessionId,
|
|
contentType: req.headers['content-type']
|
|
});
|
|
return;
|
|
}
|
|
await redisClient.publish(redisChannelForSession(sessionId), body);
|
|
res.status(202).end();
|
|
}
|