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" });
+}
@@ -0,0 +1,63 @@
+import type { Response } from "express";
+import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
+
+const DISCONNECT_GRACE_MS = 60_000;
+const IDLE_TTL_MS = 2 * 60 * 60_000;
+const MAX_AGE_MS = 24 * 60 * 60_000;
+
+/** Only transport metadata lives here; never approval arguments or results. */
+export class McpSessionLifecycle {
+  private readonly createdAt: number;
+  private lastSeenAt: number;
+  private disconnectedAt?: number;
+  private openStreams = 0;
+  private activeRequests = 0;
+
+  constructor(private readonly now: () => number = Date.now) {
+    this.createdAt = this.lastSeenAt = now();
+  }
+
+  /**
+   * Install before registering tools, using the SDK's public handler API.
+   * Track actual handler completion, not HTTP close or JSON-RPC replies:
+   * cancelled requests may finish upstream without ever sending a reply.
+   */
+  trackRequests(server: Server): void {
+    const setRequestHandler = server.setRequestHandler.bind(server);
+    server.setRequestHandler = (schema, handler) => {
+      setRequestHandler(schema, async (request, extra) => {
+        this.activeRequests += 1;
+        try { return await handler(request, extra); }
+        finally {
+          this.activeRequests -= 1;
+          this.touch();
+        }
+      });
+    };
+  }
+
+  touch(): void { this.lastSeenAt = this.now(); }
+
+  trackStream(res: Response): void {
+    this.openStreams += 1;
+    res.once("close", () => {
+      this.openStreams -= 1;
+      // A rejected duplicate GET (409) must not start the disconnect timer.
+      if (res.statusCode === 200) {
+        this.disconnectedAt = this.now();
+      }
+    });
+  }
+
+  expired(): boolean {
+    if (this.activeRequests > 0) return false;
+    const now = this.now();
+    if (now - this.createdAt >= MAX_AGE_MS) return true;
+    if (this.openStreams > 0) return false;
+    if (this.disconnectedAt !== undefined
+      && now - Math.max(this.disconnectedAt, this.lastSeenAt) >= DISCONNECT_GRACE_MS) return true;
+    // GET streams are optional in Streamable HTTP. POST-only clients retain
+    // the existing idle timeout instead of being mistaken for dead clients.
+    return now - this.lastSeenAt >= IDLE_TTL_MS;
+  }
+}
@@ -1,13 +1,11 @@
-import { randomUUID } from "node:crypto";
+import type { McpSessionLifecycle } from "./mcp-session-lifecycle.js";
+import { installMcpTransports } from "./mcp-http.js";
 import express, { type NextFunction, type Request, type Response } from "express";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
-import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
 import { hostHeaderValidation } from "@modelcontextprotocol/sdk/server/middleware/hostHeaderValidation.js";
 import { getOAuthProtectedResourceMetadataUrl, mcpAuthRouter } from "@modelcontextprotocol/sdk/server/auth/router.js";
 import { requireBearerAuth } from "@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js";
 import type { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
-import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import { z } from "zod";
 import {
  APPROVAL_BATCH_LIMIT,
@@ -28,15 +26,6 @@ import { InvoiceDeduplicationService, InvoiceDedupStore, invoiceDedupDatabasePat
 import { APPLICATION_VERSION } from "./version.js";
 import { SignedReferenceCodec } from "./signed-reference.js";

-type TransportEntry =
-  | { kind: "streamable"; transport: StreamableHTTPServerTransport; userId: string; createdAt: number; lastSeenAt: number }
-  | { kind: "sse"; transport: SSEServerTransport; userId: string; createdAt: number; lastSeenAt: number };
-
-const MCP_SESSION_IDLE_TTL_MS = 2 * 60 * 60_000;
-const MCP_SESSION_MAX_AGE_MS = 24 * 60 * 60_000;
-const MCP_SESSION_MAX_PER_USER = 8;
-const MCP_SESSION_MAX_TOTAL = 1000;
-
 export const APPROVAL_MCP_INSTRUCTIONS = [
  "先调用 approval_list_todos 获取当前待办。薪福通待办不按单据类型过滤；availableActions=[] 且 manualAction.reason=unsupported_xft_task_type 的项目仍要展示给用户，但只能引导到薪福通客户端查看和处理，禁止尝试 MCP 写入。用户只是在查看时，仅回答标题、申请人、金额、时间等业务信息，不主动讨论审批接口、批处理方式或索取写操作授权。用户明确要求处理后，在其授权范围内内部选择整组引用或自动分批；技术分批不得造成重复确认。",
  "approvalId、selectionId、targetId、单次条数限制、分批方式以及内部引用是否返回均属于 Agent 与 MCP 的内部调用细节，绝不向用户展示、解释或要求用户操作。面向用户只说明目标审批、业务动作、成功结果和需要用户决策的业务失败；不要说引用是否签发、是否失效、是否存在一键全选或如何分组。",
@@ -61,11 +50,13 @@ export function createApprovalMcpServer(
  authenticatedUser?: AuthenticatedDingTalkUser,
  audit?: (userId: string, action: string, reference: string | undefined, outcome: string) => void,
  invoiceDedup?: InvoiceDeduplicationService,
+  lifecycle?: McpSessionLifecycle,
 ): McpServer {
  const server = new McpServer(
    { name: "dingtalk-approval-mcp", version: APPLICATION_VERSION },
    { instructions: APPROVAL_MCP_INSTRUCTIONS },
  );
+  lifecycle?.trackRequests(server.server);
  const unified = new UnifiedApprovalService(config, service, xft, xftChallenges, officialOaReturns, invoiceDedup);
  const user = authenticatedUser ?? config.trustedLocalUser;
  if (!user) throw new Error("MCP 会话缺少已认证用户");
@@ -469,114 +460,15 @@ export async function startHttpServer(config: AppConfig): Promise<void> {
    ...(authMiddleware ? [authMiddleware, attachmentBody, attachmentUploadHandler] : [attachmentBody, attachmentUploadHandler]),
  );

-  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 = () => {
-    const now = Date.now();
-    for (const [id, entry] of transports) {
-      if (entry.lastSeenAt < now - MCP_SESSION_IDLE_TTL_MS || entry.createdAt < now - MCP_SESSION_MAX_AGE_MS) {
-        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, 10 * 60_000);
-  transportCleanupTimer.unref();
-
-  const mcpHandler = async (req: Request, res: Response) => {
-    let reservedUserId: string | undefined;
-    try {
-      cleanupTransports();
-      const user = requestUser(config, oauth, req.auth);
-      const sessionId = headerString(req.headers["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.lastSeenAt = Date.now();
-      if (!entry && req.method === "POST" && !sessionId && isInitializeRequest(req.body)) {
-        if (!sessionCapacityAvailable(user.userId)) {
-          res.status(429).json({ error: "MCP session limit reached; close an old Agent session and retry" });
-          return;
-        }
-        reserveSession(user.userId);
-        reservedUserId = user.userId;
-        const createdAt = Date.now();
-        const transport: StreamableHTTPServerTransport = new StreamableHTTPServerTransport({
-          sessionIdGenerator: () => randomUUID(),
-          onsessioninitialized: (id) => {
-            transports.set(id, { kind: "streamable", transport, userId: user.userId, createdAt, lastSeenAt: Date.now() });
-          },
-        });
-        transport.onclose = () => { if (transport.sessionId) transports.delete(transport.sessionId); };
-        entry = { kind: "streamable", transport, userId: user.userId, createdAt, lastSeenAt: createdAt };
-        await createApprovalMcpServer(
-          config, service, xft, undefined, officialOaReturns, user,
-          oauth ? (userId, action, reference, outcome) => oauth.audit(userId, action, reference, outcome) : undefined,
-          invoiceDedup,
-        ).connect(transport);
-      }
-      if (!entry || entry.kind !== "streamable") {
-        res.status(sessionId ? 404 : 400).json({ error: "Missing or invalid MCP session" });
-        return;
-      }
-      await entry.transport.handleRequest(req, res, req.body);
-    } catch (error) { sendServerError(res, error); }
-    finally { if (reservedUserId) releaseSessionReservation(reservedUserId); }
-  };
-  app.all("/mcp", ...(authMiddleware ? [authMiddleware, mcpHandler] : [mcpHandler]));
-
-  const sseHandler = async (req: Request, res: Response) => {
-    const user = requestUser(config, oauth, req.auth);
-    cleanupTransports();
-    if (!sessionCapacityAvailable(user.userId)) {
-      res.status(429).send("MCP session limit reached");
-      return;
-    }
-    const transport = new SSEServerTransport("/messages", res);
-    const createdAt = Date.now();
-    transports.set(transport.sessionId, { kind: "sse", transport, userId: user.userId, createdAt, lastSeenAt: createdAt });
-    res.on("close", () => transports.delete(transport.sessionId));
-    await createApprovalMcpServer(config, service, xft, undefined, officialOaReturns, user, undefined, invoiceDedup).connect(transport);
-  };
-  app.get("/sse", ...(authMiddleware ? [authMiddleware, sseHandler] : [sseHandler]));
-  const messagesHandler = async (req: Request, res: Response) => {
-    const user = requestUser(config, oauth, req.auth);
-    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.lastSeenAt = Date.now();
-    await entry.transport.handlePostMessage(req, res, req.body);
-  };
-  app.post("/messages", ...(authMiddleware ? [authMiddleware, messagesHandler] : [messagesHandler]));
+  installMcpTransports(app, {
+    authenticate: authMiddleware,
+    requestUser: (req) => requestUser(config, oauth, req.auth),
+    createServer: (user, lifecycle) => createApprovalMcpServer(
+      config, service, xft, undefined, officialOaReturns, user,
+      oauth ? (userId, action, reference, outcome) => oauth.audit(userId, action, reference, outcome) : undefined,
+      invoiceDedup, lifecycle,
+    ),
+  });
  app.get("/healthz", (_req, res) => res.json({ status: "ok", mode: config.appMode }));

  await new Promise<void>((resolve, reject) => {