UHD_Media_Tool/apps/server/src/modules/mcp/sessions.ts

// @ts-types="@modelcontextprotocol/sdk/server/index"
import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
// @ts-types="@modelcontextprotocol/sdk/server/webStandardStreamableHttp"
import type { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
import type { McpSessionInfo } from "@umt/contracts";

// MCP 会话表(进程内):会话与建立它的 API Key 绑定,换 Key 复用会话一律拒绝。
// 空闲超时回收;每个 Key 的会话数与每分钟工具调用数有上限。

export const MCP_LIMITS = {
  /** 每个 API Key 每分钟的工具调用上限 */
  toolCallsPerMinute: 60,
  /** 每个 API Key 每分钟的 HTTP 请求上限(含 list / resources / 推送流) */
  requestsPerMinute: 600,
  /** 每个 API Key 同时保持的会话上限 */
  maxSessionsPerKey: 8,
  /** 全进程会话上限 */
  maxSessions: 256,
  /** 空闲多久回收(ms) */
  idleMs: 30 * 60_000,
} as const;

export interface McpSession {
  id: string;
  transport: WebStandardStreamableHTTPServerTransport;
  server: Server;
  principalKind: string;
  principalId: string;
  principalName: string;
  clientName: string | null;
  createdAt: Date;
  lastSeenAt: Date;
  toolCalls: number;
}

/** 滑动窗口计数器:key → 最近一分钟的时间戳 */
export class RateWindow {
  private readonly hits = new Map<string, number[]>();
  constructor(
    readonly limit: number,
    readonly windowMs = 60_000,
    private readonly now: () => number = Date.now,
  ) {}

  /** 记一次;超限返回 false(不计入) */
  take(key: string): boolean {
    const t = this.now();
    const arr = (this.hits.get(key) ?? []).filter((x) => t - x < this.windowMs);
    if (arr.length >= this.limit) {
      this.hits.set(key, arr);
      return false;
    }
    arr.push(t);
    this.hits.set(key, arr);
    return true;
  }

  /** 清掉整窗口都没有记录的 key */
  sweep(): void {
    const t = this.now();
    for (const [k, arr] of this.hits) if (arr.every((x) => t - x >= this.windowMs)) this.hits.delete(k);
  }
}

export class McpSessionTable {
  readonly sessions = new Map<string, McpSession>();
  readonly toolCalls = new RateWindow(MCP_LIMITS.toolCallsPerMinute);
  readonly requests = new RateWindow(MCP_LIMITS.requestsPerMinute);
  private timer: ReturnType<typeof setInterval> | null = null;

  constructor(private readonly onClose?: (s: McpSession, reason: string) => void) {}

  countFor(principalId: string): number {
    let n = 0;
    for (const s of this.sessions.values()) if (s.principalId === principalId) n++;
    return n;
  }

  add(s: McpSession): void {
    this.sessions.set(s.id, s);
    this.ensureSweeper();
  }

  get(id: string): McpSession | undefined {
    return this.sessions.get(id);
  }

  /** 从表里摘掉并关闭传输(transport.close 会触发 onsessionclosed,重复调用无害) */
  async remove(id: string, reason: string): Promise<void> {
    const s = this.sessions.get(id);
    if (!s) return;
    this.sessions.delete(id);
    this.onClose?.(s, reason);
    await s.server.close().catch(() => {});
  }

  /** 回收空闲会话 */
  async sweep(now = Date.now()): Promise<void> {
    for (const s of [...this.sessions.values()]) {
      if (now - s.lastSeenAt.getTime() > MCP_LIMITS.idleMs) await this.remove(s.id, "idle");
    }
    this.toolCalls.sweep();
    this.requests.sweep();
  }

  async closeAll(): Promise<void> {
    if (this.timer) clearInterval(this.timer);
    this.timer = null;
    for (const id of [...this.sessions.keys()]) await this.remove(id, "shutdown");
  }

  list(): McpSessionInfo[] {
    return [...this.sessions.values()]
      .sort((a, b) => b.lastSeenAt.getTime() - a.lastSeenAt.getTime())
      .map((s) => ({
        idPrefix: s.id.slice(0, 8),
        apiKeyId: s.principalId,
        apiKeyName: s.principalName,
        clientName: s.clientName,
        createdAt: s.createdAt.toISOString(),
        lastSeenAt: s.lastSeenAt.toISOString(),
        toolCalls: s.toolCalls,
      }));
  }

  private ensureSweeper() {
    if (this.timer) return;
    this.timer = setInterval(() => void this.sweep(), 60_000);
    // 不因为回收定时器挂住进程退出
    (this.timer as unknown as { unref?: () => void }).unref?.();
    if (typeof Deno !== "undefined") Deno.unrefTimer(this.timer as unknown as number);
  }
}