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

// MCP Gateway：服务进程内的 Streamable HTTP MCP 服务（/mcp），两个引擎共用同一套工具（ADR-0006）。
// 提供内置工具（builtin.ts），并转接 Profile 勾选的第三方 MCP 服务（ADR-0021）：工具名为 <服务名>__<工具名>，
// 调用前按工具策略与权限模式审批，结果过长时截断并把完整内容存入工作区
import crypto from 'node:crypto'
import fs from 'node:fs'
import path from 'node:path'
import type { FastifyInstance } from 'fastify'
import { Server } from '@modelcontextprotocol/sdk/server/index.js'
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
import { CallToolRequestSchema, ListToolsRequestSchema, type Tool } from '@modelcontextprotocol/sdk/types.js'
import { z } from 'zod'
import type { AppSettings } from '@agent-server/shared'
import type { McpCachedTool } from '../db/schema'
import { defaultPolicy, type McpServerService } from '../services/mcp-servers'
import { BUILTIN_TOOLS, fail, type ToolContent, type ToolResult } from './builtin'
import type { GatewayRunContext, RunTokenRegistry } from './tokens'
import { isTransient, type McpServerRow, type McpUpstream } from './upstream'

export const GATEWAY_SERVER = 'gateway'
export const GATEWAY_TOOLS = BUILTIN_TOOLS.map(t => t.name)
/** 第三方工具名中服务名与工具名的分隔符 */
export const MCP_SEP = '__'

const MAX_TEXT = 100_000
const MAX_INLINE_BINARY = 5 * 1048576
const CALL_TIMEOUT_MS = 10 * 60_000
const LIST_TIMEOUT_MS = 20_000

/** 引擎等待网关工具的时间（秒）：交互超时 + 第三方工具最长执行时间 + 余量 */
export const gatewayToolTimeoutSec = (s: AppSettings) => s.limits.interactionTimeoutSec + CALL_TIMEOUT_MS / 1000 + 60

export interface GatewayDeps {
  mcp: McpServerService
  upstream: McpUpstream
}

/** 会话内"本会话允许"过的第三方工具 */
const sessionAllowed = new Map<string, Set<string>>()

function builtinTools(): Tool[] {
  return BUILTIN_TOOLS.map(t => ({
    name: t.name, description: t.description,
    inputSchema: z.toJSONSchema(t.schema) as Tool['inputSchema'],
    ...(t.readOnly ? { annotations: { readOnlyHint: true } } : {}),
  }))
}

/** 这次对话可用的第三方工具（只读模式下只提供只读工具） */
async function externalTools(ctx: GatewayRunContext, deps: GatewayDeps): Promise<{ row: McpServerRow; tool: McpCachedTool; exposed: string }[]> {
  const out: { row: McpServerRow; tool: McpCachedTool; exposed: string }[] = []
  for (const row of deps.mcp.rows(ctx.mcpServerIds)) {
    let tools = row.toolsCache
    if (!tools.length) {
      tools = await Promise.race([
        deps.mcp.listToolsFor(row).catch(() => [] as McpCachedTool[]),
        new Promise<McpCachedTool[]>(r => setTimeout(() => r([]), LIST_TIMEOUT_MS)),
      ])
    }
    for (const tool of tools) {
      const policy = row.toolPolicies[tool.name] ?? defaultPolicy(tool)
      if (policy === 'off' || (ctx.permissionMode === 'readonly' && !tool.readOnly)) continue
      out.push({ row, tool, exposed: `${row.name}${MCP_SEP}${tool.name}`.slice(0, 64) })
    }
  }
  return out
}

function buildServer(ctx: GatewayRunContext, deps: GatewayDeps): Server {
  const server = new Server({ name: 'agent-server-gateway', version: '0.1.0' }, { capabilities: { tools: {} } })

  server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [
      ...builtinTools(),
      ...(await externalTools(ctx, deps)).map(({ row, tool, exposed }) => ({
        name: exposed,
        description: `[${row.title}] ${tool.description}`.slice(0, 2000),
        inputSchema: (tool.inputSchema?.type === 'object' ? tool.inputSchema : { type: 'object', properties: {} }) as Tool['inputSchema'],
        ...(tool.readOnly ? { annotations: { readOnlyHint: true } } : {}),
      })),
    ],
  }))

  server.setRequestHandler(CallToolRequestSchema, async (req, extra) => {
    const name = req.params.name
    const args = (req.params.arguments ?? {}) as Record<string, unknown>
    const started = Date.now()
    const audit = (summary: string, detail: Record<string, unknown> = {}) =>
      ctx.audit('gateway_call', summary, { args, ms: Date.now() - started, engine: ctx.provider, ...detail })

    const builtin = BUILTIN_TOOLS.find(t => t.name === name)
    if (builtin) {
      const parsed = builtin.schema.safeParse(args)
      if (!parsed.success) return fail(`参数不正确：${parsed.error.issues.map(i => `${i.path.join('.')} ${i.message}`).join('；')}`)
      try {
        const r = await builtin.handler(parsed.data as Record<string, unknown>, ctx)
        audit(`${name}${r.isError ? '（失败）' : ''}`)
        return r
      } catch (err) {
        audit(`${name}（异常）`, { error: (err as Error).message })
        return fail((err as Error).message)
      }
    }

    const found = (await externalTools(ctx, deps)).find(t => t.exposed === name)
    if (!found) return fail(`工具不存在或当前不可用：${name}`)
    const { row, tool } = found
    const policy = row.toolPolicies[tool.name] ?? defaultPolicy(tool)
    // 需要审批：跟随权限模式（需审批时询问；自动模式直接调用；只读模式已在列表中排除非只读工具）
    if (policy === 'ask' && ctx.permissionMode === 'ask' && !sessionAllowed.get(ctx.sessionId)?.has(name)) {
      const outcome = await ctx.interact({
        kind: 'approval', tool: `mcp__${GATEWAY_SERVER}__${name}`, title: `${row.title} · ${tool.name}`,
        input: args, detail: { reason: tool.description.slice(0, 300) || undefined }, options: ['allow_once', 'allow_session', 'deny'],
      })
      ctx.audit('approval', `MCP 工具 ${name}：${outcome.resolution}`, { args })
      if (outcome.resolution !== 'allowed') return fail(`用户拒绝了这次调用${outcome.resolution === 'denied' && outcome.reason ? `：${outcome.reason}` : ''}`)
      if (outcome.scope === 'session') sessionAllowed.set(ctx.sessionId, new Set([...(sessionAllowed.get(ctx.sessionId) ?? []), name]))
    }
    try {
      const call = () => deps.upstream.callTool(row, tool.name, args, extra.signal, CALL_TIMEOUT_MS) as Promise<{ content?: ToolContent[]; structuredContent?: unknown; isError?: boolean }>
      // 偶发的网络错误：只读工具重新连接后重试一次；其他工具不重试，避免重复执行
      const r = await call().catch(async err => {
        if (!tool.readOnly || !isTransient(err)) throw err
        await deps.upstream.close(row.id)
        return call()
      })
      const result = postProcess(r, ctx.workspaceDir, name)
      audit(`${name}${result.isError ? '（失败）' : ''}`, { server: row.name })
      return result
    } catch (err) {
      audit(`${name}（异常）`, { server: row.name, error: (err as Error).message })
      return fail(`调用 ${row.title} 失败：${(err as Error).message}`)
    }
  })

  return server
}

/** 资源文件名：取 URI 最后一段，只保留安全字符；取不到时用工具名 */
function resourceFileName(uri: unknown, tool: string): string {
  let last = ''
  if (typeof uri === 'string') {
    const noQuery = uri.split(/[?#]/)[0]!
    last = noQuery.slice(noQuery.lastIndexOf('/') + 1)
    try {
      last = decodeURIComponent(last)
    } catch {
      // 保留原文
    }
  }
  const safe = last.replace(/[^\p{L}\p{N}._-]+/gu, '_').replace(/^[._]+/, '').slice(0, 120)
  return safe || `${tool.replace(/[^\w.-]/g, '_')}.bin`
}

const formatBytes = (n: number) => (n >= 1048576 ? `${(n / 1048576).toFixed(1)} MB` : n >= 1024 ? `${(n / 1024).toFixed(1)} KB` : `${n} B`)

/**
 * 只保留引擎能处理的内容；过长的文本和过大的二进制存入工作区 .agent/mcp/ 并返回路径。
 * resource（嵌入资源，如导出的 CSV、文件）一律存入工作区，只返回路径，不放进上下文
 */
function postProcess(r: { content?: ToolContent[]; structuredContent?: unknown; isError?: boolean }, workspaceDir: string, tool: string): ToolResult {
  const dir = path.join(workspaceDir, '.agent', 'mcp')
  const save = (ext: string, data: string | Buffer) => {
    fs.mkdirSync(dir, { recursive: true })
    const file = path.join(dir, `${tool.replace(/[^\w.-]/g, '_')}-${Date.now()}-${crypto.randomBytes(3).toString('hex')}.${ext}`)
    fs.writeFileSync(file, data)
    return path.relative(workspaceDir, file)
  }
  /** 按资源原名保存，重名时加序号 */
  const saveNamed = (name: string, data: string | Buffer) => {
    fs.mkdirSync(dir, { recursive: true })
    const ext = path.extname(name)
    const stem = name.slice(0, name.length - ext.length)
    let file = path.join(dir, name)
    for (let i = 2; fs.existsSync(file); i++) file = path.join(dir, `${stem}-${i}${ext}`)
    fs.writeFileSync(file, data)
    return path.relative(workspaceDir, file)
  }
  const content: ToolContent[] = []
  for (const c of r.content ?? []) {
    if (c.type === 'resource' && c.resource && typeof c.resource === 'object') {
      const res = c.resource as { uri?: unknown; mimeType?: unknown; text?: unknown; blob?: unknown }
      const data = typeof res.blob === 'string' ? Buffer.from(res.blob, 'base64') : typeof res.text === 'string' ? res.text : null
      if (data === null) {
        content.push({ type: 'text', text: '（资源没有内容）' })
        continue
      }
      const saved = saveNamed(resourceFileName(res.uri, tool), data)
      const size = typeof data === 'string' ? Buffer.byteLength(data) : data.length
      const mime = typeof res.mimeType === 'string' && res.mimeType ? res.mimeType : typeof data === 'string' ? 'text/plain' : 'application/octet-stream'
      content.push({ type: 'text', text: `已保存到 ${saved}（${formatBytes(size)}，${mime}）` })
    } else if (c.type === 'text' && typeof c.text === 'string' && c.text.length > MAX_TEXT) {
      const saved = save('txt', c.text)
      content.push({ type: 'text', text: `${c.text.slice(0, MAX_TEXT)}\n\n（内容过长，以上为前一部分；完整内容已保存到 ${saved}）` })
    } else if ((c.type === 'image' || c.type === 'audio') && typeof c.data === 'string' && c.data.length > MAX_INLINE_BINARY) {
      const ext = String(c.mimeType ?? '').split('/')[1]?.replace(/[^\w]/g, '') || 'bin'
      content.push({ type: 'text', text: `（${c.type === 'image' ? '图片' : '音频'}过大，已保存到 ${save(ext, Buffer.from(c.data, 'base64'))}）` })
    } else {
      content.push(c)
    }
  }
  if (!content.length && r.structuredContent !== undefined) content.push({ type: 'text', text: JSON.stringify(r.structuredContent, null, 2).slice(0, MAX_TEXT) })
  if (!content.length) content.push({ type: 'text', text: '（无返回内容）' })
  return { content, ...(r.isError ? { isError: true } : {}) }
}

export function registerGateway(f: FastifyInstance, tokens: RunTokenRegistry, deps: GatewayDeps) {
  f.post('/mcp', async (req, reply) => {
    const auth = req.headers.authorization
    const ctx = tokens.get(auth?.startsWith('Bearer ') ? auth.slice(7).trim() : undefined)
    if (!ctx) return reply.code(401).send({ jsonrpc: '2.0', error: { code: -32001, message: '无效或已过期的运行令牌' }, id: null })
    // 无状态模式：每个请求一个 server/transport 实例，绑定到该 Run 的上下文
    const server = buildServer(ctx, deps)
    const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined })
    reply.hijack()
    reply.raw.on('close', () => {
      void transport.close()
      void server.close()
    })
    await server.connect(transport)
    await transport.handleRequest(req.raw, reply.raw, req.body)
  })
  const notAllowed = { jsonrpc: '2.0', error: { code: -32000, message: 'Method not allowed' }, id: null }
  f.get('/mcp', async (_req, reply) => reply.code(405).send(notAllowed))
  f.delete('/mcp', async (_req, reply) => reply.code(405).send(notAllowed))
}

/** 会话删除时清除"本会话允许" */
export function forgetGatewaySession(sessionId: string) {
  sessionAllowed.delete(sessionId)
}