AgentServer/apps/server/test/gateway.test.ts

import fs from 'node:fs'
import http from 'node:http'
import path from 'node:path'
import type { AddressInfo } from 'node:net'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import pino from 'pino'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { createApp, type App } from '../src/app'
import { buildServer } from '../src/http/server'
import type { GatewayRunContext } from '../src/mcp-gateway/tokens'
import { resolveAccess } from '../src/services/access-policy'
import { config } from '../src/config'

let app: App
let server: Awaited<ReturnType<typeof buildServer>>
let base: string
let site: http.Server
let siteUrl: string
const audits: string[] = []

function runCtx(over: Partial<GatewayRunContext> = {}): GatewayRunContext {
  const ws = path.join(config.dataDir, 'workspaces', `ses_gw_${Math.random().toString(36).slice(2)}`)
  fs.mkdirSync(path.join(ws, 'outputs'), { recursive: true })
  const settings = app.settings.get()
  return {
    sessionId: 's', runId: 'r', provider: 'codex', workspaceDir: ws, permissionMode: 'ask', settings,
    access: resolveAccess({ policy: { mode: 'restricted', read: [], write: [], deny: [] }, knowledgeRoots: [], selectedKnowledgeRootIds: [], workspaceDir: ws, settings }),
    knowledgeRoots: [{ id: 'kb', name: '账单', path: '/data/bills', description: '家庭账单', access: 'read', enabled: true, check: { exists: true, readable: true, isDirectory: true, denied: null, underHome: false, message: null } }],
    mcpServerIds: [],
    audit: (kind, summary) => audits.push(`${kind}:${summary}`),
    interact: async () => ({ resolution: 'allowed', scope: 'once' }),
    ...over,
  }
}

async function connect(token: string) {
  const client = new Client({ name: 'test', version: '0' })
  await client.connect(new StreamableHTTPClientTransport(new URL(`${base}/mcp`), { requestInit: { headers: { Authorization: `Bearer ${token}` } } }))
  return client
}
const text = (r: Awaited<ReturnType<Client['callTool']>>) => (r.content as { text: string }[]).map(c => c.text).join('\n')

beforeAll(async () => {
  const log = pino({ level: 'silent' })
  app = createApp(log, { providers: {} })
  server = await buildServer(app, { logger: log })
  await server.listen({ host: '127.0.0.1', port: 0 })
  base = `http://127.0.0.1:${(server.server.address() as AddressInfo).port}`
  site = http.createServer((_req, res) => {
    res.setHeader('content-type', 'text/html; charset=utf-8')
    res.end('<html><head><script>evil()</script></head><body><h1>标题</h1><p>正文 <b>加粗</b></p></body></html>')
  })
  await new Promise<void>(r => site.listen(0, '127.0.0.1', r))
  siteUrl = `http://localhost:${(site.address() as AddressInfo).port}/page`
})
afterAll(async () => {
  await server.close()
  site.close()
})

describe('MCP Gateway', () => {
  it('无效令牌被拒绝', async () => {
    const res = await fetch(`${base}/mcp`, { method: 'POST', headers: { authorization: 'Bearer nope', 'content-type': 'application/json', accept: 'application/json, text/event-stream' }, body: '{}' })
    expect(res.status).toBe(401)
    await expect(connect('nope')).rejects.toThrow()
  })

  it('列出工具；资料库列表；文档转换写入工作区', async () => {
    const ctx = runCtx()
    const client = await connect(app.gatewayTokens.issue(ctx))
    const tools = (await client.listTools()).tools.map(t => t.name).sort()
    expect(tools).toEqual(['convert_document', 'kb_list_roots', 'recalc_workbook', 'web_fetch'])
    expect(text(await client.callTool({ name: 'kb_list_roots', arguments: {} }))).toContain('账单：/data/bills（只读）')

    fs.writeFileSync(path.join(ctx.workspaceDir, 'a.txt'), '你好')
    const r = await client.callTool({ name: 'convert_document', arguments: { input: 'a.txt', format: 'pdf' } })
    expect(text(r)).toBe('已生成：outputs/a.pdf')
    expect(fs.existsSync(path.join(ctx.workspaceDir, 'outputs/a.pdf'))).toBe(true)
    expect(audits.some(a => a === 'gateway_call:convert_document')).toBe(true)
    await client.close()
  }, 60_000)

  it('只读模式不能生成文件；工作区外的输入被拒绝', async () => {
    const ctx = runCtx({ permissionMode: 'readonly' })
    const client = await connect(app.gatewayTokens.issue(ctx))
    const r = await client.callTool({ name: 'convert_document', arguments: { input: 'a.txt', format: 'pdf' } })
    expect(r.isError).toBe(true)
    expect(text(r)).toContain('只读模式')
    const ctx2 = runCtx()
    const c2 = await connect(app.gatewayTokens.issue(ctx2))
    const r2 = await c2.callTool({ name: 'convert_document', arguments: { input: '/etc/hostname', format: 'pdf' } })
    expect(text(r2)).toContain('工作区内')
    await client.close()
    await c2.close()
  })

  it('web_fetch 拦截内网地址（包括解析到内网的域名），放开后能转成 Markdown', async () => {
    const ctx = runCtx()
    const client = await connect(app.gatewayTokens.issue(ctx))
    const blocked = await client.callTool({ name: 'web_fetch', arguments: { url: siteUrl } })
    expect(blocked.isError).toBe(true)
    expect(text(blocked)).toMatch(/不允许访问内网地址：localhost/)
    expect(audits.some(a => a.startsWith('network_denied:'))).toBe(true)
    await client.close()

    const open = runCtx({ settings: { ...app.settings.get(), network: { ...app.settings.get().network, blockPrivateFetch: false } } })
    const c2 = await connect(app.gatewayTokens.issue(open))
    const r = text(await c2.callTool({ name: 'web_fetch', arguments: { url: siteUrl } }))
    expect(r).toContain('# 标题')
    expect(r).toContain('正文 **加粗**')
    expect(r).not.toContain('evil')
    await c2.close()
  })

  it('令牌作废后不能再调用', async () => {
    const token = app.gatewayTokens.issue(runCtx())
    app.gatewayTokens.revoke(token)
    await expect(connect(token)).rejects.toThrow()
  })

  it('Codex 会话使用固定令牌：每次 Run 重新绑定，Run 之间失效', () => {
    const reg = app.gatewayTokens
    const t1 = reg.issue(runCtx({ sessionId: 'ses_fixed', runId: 'r1' }), { stablePerSession: true })
    reg.revoke(t1)
    expect(reg.get(t1)).toBeNull()
    const t2 = reg.issue(runCtx({ sessionId: 'ses_fixed', runId: 'r2' }), { stablePerSession: true })
    expect(t2).toBe(t1)
    expect(reg.get(t2)?.runId).toBe('r2')
    // 其他会话、非固定令牌各不相同
    expect(reg.issue(runCtx({ sessionId: 'ses_other' }), { stablePerSession: true })).not.toBe(t1)
    expect(reg.issue(runCtx({ sessionId: 'ses_fixed' }))).not.toBe(t1)
    reg.forgetSession('ses_fixed')
    expect(reg.get(t1)).toBeNull()
    expect(reg.issue(runCtx({ sessionId: 'ses_fixed' }), { stablePerSession: true })).not.toBe(t1)
  })
})