AgentServer/apps/server/src/mcp-gateway/upstream.ts

// 第三方 MCP 服务的连接管理（ADR-0021）：MCP Gateway 作为客户端连接它们，把工具转接给两个引擎。
// - http：Streamable HTTP（失败时退回 SSE），请求头可携带令牌
// - stdio：本地进程，在 bubblewrap 沙箱中运行：系统目录只读、独立的可写 home、按配置开放网络与额外路径；
//   宿主 Unix socket（docker.sock 等）不可达；密钥只通过环境变量传给该进程
// 连接按需建立、多个对话共用，空闲 10 分钟后关闭；配置变化时重新连接
import fs from 'node:fs'
import path from 'node:path'
import type { Readable } from 'node:stream'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import { APP_VERSION, config } from '../config'
import type { McpCachedTool, mcpServers } from '../db/schema'
import { hostSockets } from '../services/office'

export type McpServerRow = typeof mcpServers.$inferSelect

const CONNECT_TIMEOUT_MS = 120_000
const IDLE_MS = 10 * 60_000
const STDERR_LINES = 40
export const BWRAP = ['/usr/bin/bwrap', '/bin/bwrap'].find(p => fs.existsSync(p)) ?? null

interface Conn {
  key: string
  client: Client
  stderr: string[]
  lastUsed: number
  /** 正在执行的调用数：有调用时不因空闲关闭（后台调用可能执行数小时） */
  inFlight: number
}

/** 影响连接的配置；变化时需要重新连接 */
const connKey = (s: McpServerRow) => JSON.stringify([s.transport, s.url, s.headers, s.command, s.args, s.env, s.allowNetwork, s.readPaths, s.writePaths])

/** 本地进程的独立 home（npx / uvx 缓存等），在数据目录下 */
export const mcpHomeDir = (id: string) => path.join(config.dataDir, 'data', 'mcp', id)

export function stdioSandboxArgs(s: Pick<McpServerRow, 'allowNetwork' | 'readPaths' | 'writePaths'>, home: string): string[] {
  return [
    '--ro-bind', '/usr', '/usr',
    '--ro-bind', '/etc', '/etc',
    '--symlink', 'usr/lib', '/lib',
    '--symlink', 'usr/lib64', '/lib64',
    '--symlink', 'usr/bin', '/bin',
    '--symlink', 'usr/sbin', '/sbin',
    ...hostSockets().flatMap(p => ['--ro-bind', '/dev/null', p]),
    '--dev', '/dev',
    '--proc', '/proc',
    '--tmpfs', '/tmp',
    '--bind', home, home,
    ...s.readPaths.flatMap(p => ['--ro-bind', p, p]),
    ...s.writePaths.flatMap(p => ['--bind', p, p]),
    '--unshare-all',
    ...(s.allowNetwork ? ['--share-net'] : []),
    '--die-with-parent',
    '--new-session',
    '--chdir', home,
    '--setenv', 'HOME', home,
    '--setenv', 'TMPDIR', '/tmp',
    '--setenv', 'LANG', 'C.UTF-8',
    '--setenv', 'PATH', '/usr/local/bin:/usr/bin:/bin',
  ]
}

/** 网络层的偶发错误（连接被重置、超时等），可以安全地重新连接 */
export const isTransient = (err: unknown) => /fetch failed|ECONNRESET|ETIMEDOUT|ECONNREFUSED|EAI_AGAIN|socket hang up|UND_ERR|Connection closed/i.test(String((err as Error)?.message ?? err) + String((err as { cause?: { code?: string } })?.cause?.code ?? ''))

function withTimeout<T>(p: Promise<T>, ms: number, what: string): Promise<T> {
  let timer: NodeJS.Timeout
  return Promise.race([p, new Promise<T>((_, reject) => { timer = setTimeout(() => reject(new Error(`${what}超时（${Math.round(ms / 1000)} 秒）`)), ms) })])
    .finally(() => clearTimeout(timer))
}

export class McpUpstream {
  private conns = new Map<string, Conn>()
  private pending = new Map<string, Promise<Conn>>()
  private errors = new Map<string, { error: string; stderr: string | null }>()
  private sweeper: NodeJS.Timeout

  constructor() {
    this.sweeper = setInterval(() => this.closeIdle(), 60_000)
    this.sweeper.unref()
  }

  status(id: string): { status: 'connected' | 'idle' | 'error'; lastError: string | null; stderrTail: string | null } {
    const c = this.conns.get(id)
    const e = this.errors.get(id)
    if (c) return { status: 'connected', lastError: null, stderrTail: c.stderr.length ? c.stderr.join('\n') : null }
    return { status: e ? 'error' : 'idle', lastError: e?.error ?? null, stderrTail: e?.stderr ?? null }
  }

  /** 列出工具（包括分页） */
  async listTools(s: McpServerRow): Promise<McpCachedTool[]> {
    const c = await this.conn(s)
    return listAll(c.client)
  }

  async callTool(s: McpServerRow, name: string, args: Record<string, unknown>, signal: AbortSignal, timeoutMs: number) {
    const c = await this.conn(s)
    c.lastUsed = Date.now()
    c.inFlight++
    try {
      return await c.client.callTool({ name, arguments: args }, undefined, { signal, timeout: timeoutMs, resetTimeoutOnProgress: true, maxTotalTimeout: timeoutMs * 2 })
    } finally {
      c.inFlight--
      c.lastUsed = Date.now()
    }
  }

  /** 用一份配置临时连接并列出工具（后台"测试连接"），不影响正在使用的连接 */
  async test(s: McpServerRow): Promise<{ tools: McpCachedTool[]; stderr: string[] }> {
    const stderr: string[] = []
    const client = await this.connect(s, stderr)
    try {
      return { tools: await listAll(client), stderr }
    } finally {
      await client.close().catch(() => {})
    }
  }

  async close(id: string) {
    const c = this.conns.get(id)
    this.conns.delete(id)
    this.errors.delete(id)
    await c?.client.close().catch(() => {})
  }

  async closeAll() {
    clearInterval(this.sweeper)
    await Promise.all([...this.conns.keys()].map(id => this.close(id)))
  }

  private closeIdle() {
    for (const [id, c] of this.conns) if (c.inFlight === 0 && Date.now() - c.lastUsed > IDLE_MS) void this.close(id)
  }

  private async conn(s: McpServerRow): Promise<Conn> {
    const key = connKey(s)
    const existing = this.conns.get(s.id)
    if (existing && existing.key === key) {
      existing.lastUsed = Date.now()
      return existing
    }
    if (existing) await this.close(s.id)
    const inflight = this.pending.get(s.id)
    if (inflight) return inflight
    const stderr: string[] = []
    // 远程服务偶发的网络错误重试一次
    const connectOnce = () => this.connect(s, stderr)
    const p = connectOnce().catch(err => (s.transport === 'http' && isTransient(err) ? connectOnce() : Promise.reject(err))).then(client => {
      const c: Conn = { key, client, stderr, lastUsed: Date.now(), inFlight: 0 }
      this.conns.set(s.id, c)
      this.errors.delete(s.id)
      // 连接断开（进程退出、服务端关闭）后下次调用重新连接
      client.onclose = () => {
        if (this.conns.get(s.id) === c) this.conns.delete(s.id)
      }
      return c
    }, (err: Error) => {
      this.errors.set(s.id, { error: err.message, stderr: stderr.length ? stderr.join('\n') : null })
      throw err
    }).finally(() => this.pending.delete(s.id))
    this.pending.set(s.id, p)
    return p
  }

  private async connect(s: McpServerRow, stderr: string[]): Promise<Client> {
    const client = new Client({ name: 'agent-server-gateway', version: APP_VERSION })
    if (s.transport === 'http') {
      const url = new URL(s.url)
      const requestInit: RequestInit = { headers: s.headers }
      try {
        await withTimeout(client.connect(new StreamableHTTPClientTransport(url, { requestInit })), CONNECT_TIMEOUT_MS, '连接')
      } catch (err) {
        // 旧版服务只支持 SSE
        const fallback = new Client({ name: 'agent-server-gateway', version: APP_VERSION })
        try {
          await withTimeout(fallback.connect(new SSEClientTransport(url, { requestInit, eventSourceInit: { fetch: (u, init) => fetch(u, { ...init, headers: { ...(init?.headers as Record<string, string>), ...s.headers } }) } })), CONNECT_TIMEOUT_MS, '连接')
          return fallback
        } catch {
          throw err
        }
      }
      return client
    }
    if (!BWRAP) throw new Error('服务器上没有 bubblewrap（bwrap），无法在沙箱中运行本地 MCP 服务')
    const home = mcpHomeDir(s.id)
    fs.mkdirSync(home, { recursive: true, mode: 0o700 })
    const transport = new StdioClientTransport({
      command: BWRAP,
      args: [...stdioSandboxArgs(s, home), '--', s.command, ...s.args],
      // 密钥只放在环境变量里（不出现在命令行参数中）
      env: { ...s.env, HOME: home },
      stderr: 'pipe',
      cwd: home,
    })
    ;(transport.stderr as Readable | null)?.on('data', (d: Buffer) => {
      for (const line of d.toString('utf8').split('\n')) if (line.trim()) stderr.push(line.slice(0, 500))
      if (stderr.length > STDERR_LINES) stderr.splice(0, stderr.length - STDERR_LINES)
    })
    try {
      await withTimeout(client.connect(transport), CONNECT_TIMEOUT_MS, '启动')
    } catch (err) {
      await transport.close().catch(() => {})
      const tail = stderr.slice(-5).join('\n')
      throw new Error(`${(err as Error).message}${tail ? `\n${tail}` : ''}`)
    }
    return client
  }
}

async function listAll(client: Client): Promise<McpCachedTool[]> {
  const tools: McpCachedTool[] = []
  let cursor: string | undefined
  for (let page = 0; page < 20; page++) {
    const r = await withTimeout(client.listTools(cursor ? { cursor } : undefined), 60_000, '列出工具')
    for (const t of r.tools) {
      tools.push({ name: t.name, description: t.description ?? '', inputSchema: t.inputSchema as Record<string, unknown>, readOnly: t.annotations?.readOnlyHint === true })
    }
    cursor = r.nextCursor
    if (!cursor) break
  }
  return tools
}