fix: reclaim disconnected gateway MCP sessions

2026-10-28 03:05:10 +00:00
parent 430d45014b
commit 2213ce7570
5 changed files with 465 additions and 121 deletions
@@ -0,0 +1,148 @@
+import { randomUUID } from "node:crypto";
+import type { Express, Request, Response, RequestHandler } from "express";
+import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
+import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
+import type { AuthenticatedDingTalkUser } from "./xft.js";
+import { McpSessionLifecycle } from "./mcp-session-lifecycle.js";
+
+type TransportEntry =
+  | { kind: "streamable"; transport: StreamableHTTPServerTransport; userId: string; lifecycle: McpSessionLifecycle }
+  | { kind: "sse"; transport: SSEServerTransport; userId: string; lifecycle: McpSessionLifecycle };
+
+const MCP_SESSION_MAX_PER_USER = 8;
+const MCP_SESSION_MAX_TOTAL = 1000;
+
+interface McpTransportOptions {
+  authenticate?: RequestHandler;
+  // Production always resolves this from verified OAuth, never request parameters.
+  requestUser: (req: Request) => AuthenticatedDingTalkUser;
+  createServer: (user: AuthenticatedDingTalkUser, lifecycle: McpSessionLifecycle) => McpServer;
+  now?: () => number;
+}
+
+export function installMcpTransports(app: Express, options: McpTransportOptions) {
+  const now = options.now ?? Date.now;
+  const transports = new Map<string, TransportEntry>();
+  let pendingTransportSessions = 0;
+  const pendingTransportSessionsByUser = new Map<string, number>();
+  const closeTransport = (id: string, entry: TransportEntry) => {
+    transports.delete(id);
+    void entry.transport.close().catch(() => undefined);
+  };
+  const cleanupTransports = () => {
+    for (const [id, entry] of transports) {
+      if (entry.lifecycle.expired()) {
+        closeTransport(id, entry);
+      }
+    }
+  };
+  const sessionCapacityAvailable = (userId: string) => transports.size + pendingTransportSessions < MCP_SESSION_MAX_TOTAL
+    && [...transports.values()].filter((entry) => entry.userId === userId).length
+      + (pendingTransportSessionsByUser.get(userId) ?? 0) < MCP_SESSION_MAX_PER_USER;
+  const reserveSession = (userId: string) => {
+    pendingTransportSessions += 1;
+    pendingTransportSessionsByUser.set(userId, (pendingTransportSessionsByUser.get(userId) ?? 0) + 1);
+  };
+  const releaseSessionReservation = (userId: string) => {
+    pendingTransportSessions = Math.max(0, pendingTransportSessions - 1);
+    const remaining = Math.max(0, (pendingTransportSessionsByUser.get(userId) ?? 1) - 1);
+    if (remaining === 0) pendingTransportSessionsByUser.delete(userId);
+    else pendingTransportSessionsByUser.set(userId, remaining);
+  };
+  const transportCleanupTimer = setInterval(cleanupTransports, 15_000);
+  transportCleanupTimer.unref();
+
+  const mcpHandler = async (req: Request, res: Response) => {
+    let reservedUserId: string | undefined;
+    try {
+      cleanupTransports();
+      const user = options.requestUser(req);
+      const sessionId = req.get("mcp-session-id");
+      let entry = sessionId ? transports.get(sessionId) : undefined;
+      if (entry && entry.kind !== "streamable") {
+        res.status(400).json({ error: "Invalid MCP session transport" });
+        return;
+      }
+      if (entry && entry.userId !== user.userId) {
+        res.status(403).json({ error: "MCP session belongs to another authenticated user" });
+        return;
+      }
+      if (entry) entry.lifecycle.touch();
+      if (!entry && req.method === "POST" && !sessionId && isInitializeRequest(req.body)) {
+        if (!sessionCapacityAvailable(user.userId)) {
+          res.set("Retry-After", "60").status(429).json({ error: "MCP session limit reached; close an old Agent session and retry" });
+          return;
+        }
+        reserveSession(user.userId);
+        reservedUserId = user.userId;
+        const transport: StreamableHTTPServerTransport = new StreamableHTTPServerTransport({
+          sessionIdGenerator: () => randomUUID(),
+          onsessioninitialized: (id) => {
+            transports.set(id, entry!);
+            // The live entry now owns the slot; do not count its reservation twice.
+            releaseSessionReservation(user.userId);
+            reservedUserId = undefined;
+          },
+        });
+        transport.onclose = () => { if (transport.sessionId) transports.delete(transport.sessionId); };
+        entry = { kind: "streamable", transport, userId: user.userId, lifecycle: new McpSessionLifecycle(now) };
+        await options.createServer(user, entry.lifecycle).connect(transport);
+      }
+      if (!entry || entry.kind !== "streamable") {
+        res.status(sessionId ? 404 : 400).json({ error: "Missing or invalid MCP session" });
+        return;
+      }
+      if (req.method === "GET") entry.lifecycle.trackStream(res);
+      await entry.transport.handleRequest(req, res, req.body);
+    } catch (error) { sendServerError(res, error); }
+    finally { if (reservedUserId) releaseSessionReservation(reservedUserId); }
+  };
+  app.all("/mcp", ...(options.authenticate ? [options.authenticate, mcpHandler] : [mcpHandler]));
+
+  const sseHandler = async (req: Request, res: Response) => {
+    const user = options.requestUser(req);
+    cleanupTransports();
+    if (!sessionCapacityAvailable(user.userId)) {
+      res.set("Retry-After", "60").status(429).send("MCP session limit reached");
+      return;
+    }
+    const transport = new SSEServerTransport("/messages", res);
+    const lifecycle = new McpSessionLifecycle(now);
+    lifecycle.trackStream(res);
+    transports.set(transport.sessionId, { kind: "sse", transport, userId: user.userId, lifecycle });
+    res.on("close", () => transports.delete(transport.sessionId));
+    await options.createServer(user, lifecycle).connect(transport);
+  };
+  app.get("/sse", ...(options.authenticate ? [options.authenticate, sseHandler] : [sseHandler]));
+  const messagesHandler = async (req: Request, res: Response) => {
+    const user = options.requestUser(req);
+    const entry = transports.get(String(req.query.sessionId ?? ""));
+    if (!entry || entry.kind !== "sse") {
+      res.status(404).send("Invalid SSE session");
+      return;
+    }
+    if (entry.userId !== user.userId) {
+      res.status(403).send("Invalid SSE session");
+      return;
+    }
+    entry.lifecycle.touch();
+    await entry.transport.handlePostMessage(req, res, req.body);
+  };
+  app.post("/messages", ...(options.authenticate ? [options.authenticate, messagesHandler] : [messagesHandler]));
+  return {
+    cleanup: cleanupTransports,
+    async close() {
+      clearInterval(transportCleanupTimer);
+      const entries = [...transports.values()];
+      transports.clear();
+      await Promise.all(entries.map((entry) => entry.transport.close()));
+    },
+  };
+}
+
+function sendServerError(res: Response, _error: unknown): void {
+  console.error("[dingtalk-approval-mcp] MCP transport request failed");
+  if (!res.headersSent) res.status(500).json({ error: "Internal server error" });
+}