refactor, improve docs

external-oauth/mcp-server/src/handlers/sse.ts

@@ -0,0 +1,100 @@
+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();
+}