AgentServer/apps/server/test/mcp-script.test.ts

// 脚本调用工具（docs/design/17-run-script.md）：模型写 Python 脚本，由网关在沙箱中执行并批量调用 MCP 工具
import fs from 'node:fs'
import http from 'node:http'
import type { AddressInfo } from 'node:net'
import path from 'node:path'
import pino from 'pino'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { z } from 'zod'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
import type { ScriptProgressPayload } from '@agent-server/shared'
import { createApp, type App } from '../src/app'
import { config } from '../src/config'
import { buildServer } from '../src/http/server'
import { BWRAP } from '../src/mcp-gateway/upstream'
import type { GatewayRunContext } from '../src/mcp-gateway/tokens'
import { resolveAccess } from '../src/services/access-policy'
import type { InteractionOutcome, NewApproval } from '../src/services/interactions'

let app: App
let server: Awaited<ReturnType<typeof buildServer>>
let base: string
let upstream: http.Server
let serverId: string
const flags: string[] = []
/** 测试的数据目录是临时目录，使用本机的 Python 运行环境（没有时跳过） */
const PY_RUNTIME = '/srv/agent-server/runtime/py'
const HAS_PY = fs.existsSync(path.join(PY_RUNTIME, 'bin/python'))

function startUpstream(): Promise<string> {
  upstream = http.createServer(async (req, res) => {
    const mcp = new McpServer({ name: 'dev', version: '1.0.0' })
    mcp.registerTool('get_item', {
      description: '读取一条数据', inputSchema: { i: z.number().int() }, outputSchema: { i: z.number(), value: z.number() }, annotations: { readOnlyHint: true },
    }, async ({ i }) => (i < 0
      ? { isError: true, content: [{ type: 'text', text: `编号 ${i} 不存在` }] }
      : { content: [{ type: 'text', text: JSON.stringify({ i, value: i * 2 }) }], structuredContent: { i, value: i * 2 } }))
    mcp.registerTool('set_flag', { description: '设置标记', inputSchema: { name: z.string() } }, async ({ name }) => {
      flags.push(name)
      return { content: [{ type: 'text', text: `ok:${name}` }] }
    })
    mcp.registerTool('wipe', { description: '清空', inputSchema: {}, annotations: { destructiveHint: true } }, async () => ({ content: [{ type: 'text', text: 'wiped' }] }))
    const t = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined })
    res.on('close', () => { void t.close(); void mcp.close() })
    await mcp.connect(t)
    let body = ''
    req.on('data', c => (body += c))
    req.on('end', () => void t.handleRequest(req, res, body ? JSON.parse(body) : undefined))
  })
  return new Promise(r => upstream.listen(0, '127.0.0.1', () => r(`http://127.0.0.1:${(upstream.address() as AddressInfo).port}/mcp`)))
}

const approvals: NewApproval[] = []
let answer: InteractionOutcome = { resolution: 'allowed', scope: 'once' }
const progress: ScriptProgressPayload[] = []
const audits: { summary: string; detail: Record<string, unknown> }[] = []
function runCtx(over: Partial<GatewayRunContext> = {}): GatewayRunContext {
  const ws = path.join(config.dataDir, 'workspaces', `ses_script_${Math.random().toString(36).slice(2)}`)
  fs.mkdirSync(ws, { recursive: true })
  const settings = structuredClone(app.settings.get())
  settings.paths.pythonRuntime = PY_RUNTIME
  return {
    sessionId: `ses_${Math.random().toString(36).slice(2)}`, runId: 'r', provider: 'claude', workspaceDir: ws, permissionMode: 'ask', settings,
    access: resolveAccess({ policy: { mode: 'restricted', read: [], write: [], deny: [] }, knowledgeRoots: [], selectedKnowledgeRootIds: [], workspaceDir: ws, settings }),
    knowledgeRoots: [], mcpServerIds: [serverId], audit: (_k, summary, detail) => void audits.push({ summary, detail: detail as Record<string, unknown> }),
    interact: async req => { approvals.push(req as NewApproval); return answer },
    progress: p => void progress.push(p),
    ...over,
  }
}
async function gateway(ctx: GatewayRunContext) {
  const client = new Client({ name: 'test', version: '1' })
  await client.connect(new StreamableHTTPClientTransport(new URL(`${base}/mcp`), { requestInit: { headers: { Authorization: `Bearer ${app.gatewayTokens.issue(ctx)}` } } }))
  return client
}
async function run(ctx: GatewayRunContext, code: string, tools: string[] = [], description = '测试脚本') {
  const c = await gateway(ctx)
  try {
    const r = await c.callTool({ name: 'run_script', arguments: { code, tools, description } }, undefined, { timeout: 120_000 }) as { content: { text: string }[]; isError?: boolean }
    return { text: r.content[0]!.text, isError: !!r.isError }
  } finally {
    await c.close()
  }
}

beforeAll(async () => {
  const log = pino({ level: 'silent' })
  app = createApp(log, { providers: {} })
  server = await buildServer(app, { logger: log })
  base = await server.listen({ port: 0, host: '127.0.0.1' })
  const url = await startUpstream()
  serverId = app.mcpServers.create({ name: 'dev', title: '设备服务', description: '', transport: 'http', enabled: true, url, headers: [], command: '', args: [], env: [], allowNetwork: false, readPaths: [], writePaths: [] }).id
  await app.mcpServers.refresh(serverId)
})
afterAll(async () => {
  app.mcpTasks.stop()
  app.backgroundCalls.stop()
  await app.mcpUpstream.closeAll()
  await server.close()
  upstream.close()
})

describe.skipIf(!BWRAP || !HAS_PY)('脚本调用工具', () => {
  it('有第三方工具时才提供 run_script', async () => {
    const none = await gateway(runCtx({ mcpServerIds: [] }))
    expect((await none.listTools()).tools.map(t => t.name)).not.toContain('run_script')
    await none.close()
    const c = await gateway(runCtx())
    const tool = (await c.listTools()).tools.find(t => t.name === 'run_script')!
    expect(tool.description).toContain('from agent_tools import')
    await c.close()
  })

  it('批量调用只读工具、在脚本中汇总并写文件，只把打印的内容交回；记录明细、审计与进度', async () => {
    const ctx = runCtx()
    progress.length = 0
    audits.length = 0
    const r = await run(ctx, `
from agent_tools import call, call_many, progress, ToolError
import pandas as pd
rows = call_many("dev__get_item", [{"i": i} for i in range(300)])
progress(300, 300, "汇总")
df = pd.DataFrame(rows)
import os
os.makedirs("outputs", exist_ok=True)
df.to_csv("outputs/items.csv", index=False)
try:
    call("dev__get_item", {"i": -1})
except ToolError as e:
    print("出错：", e.message)
print("合计", int(df.value.sum()), "条数", len(df))
`, [], '读取 300 条并汇总')
    expect(r.isError).toBe(false)
    expect(r.text).toContain('合计 89700 条数 300')
    expect(r.text).toContain('出错： 编号 -1 不存在')
    expect(r.text).toMatch(/调用工具 301 次（失败 1 次）/)
    expect(fs.readFileSync(path.join(ctx.workspaceDir, 'outputs/items.csv'), 'utf8').split('\n')[0]).toBe('i,value')
    expect(approvals).toHaveLength(0)

    const logDir = /调用明细：(\S+)\/calls\.jsonl/.exec(r.text)![1]!
    expect(fs.readFileSync(path.join(ctx.workspaceDir, logDir, 'calls.jsonl'), 'utf8').trim().split('\n')).toHaveLength(301)
    expect(fs.readFileSync(path.join(ctx.workspaceDir, logDir, 'script.py'), 'utf8')).toContain('call_many')

    expect(progress.at(-1)).toMatchObject({ title: '读取 300 条并汇总', status: 'completed', calls: 301, failures: 1, done: 300, total: 300, message: '汇总' })
    const a = audits.find(x => x.summary.startsWith('run_script'))!
    expect(a.summary).toBe('run_script：读取 300 条并汇总（调用工具 301 次，失败 1 次）')
    expect(a.detail).toMatchObject({ byTool: { dev__get_item: 301 }, exitCode: 0 })
  })

  it('写工具：没有列出的不能调用；列出的执行前审批一次；拒绝则不执行；破坏性工具不能在脚本中调用', async () => {
    const code = 'from agent_tools import call\nfor n in ["a", "b", "c"]:\n    print(call("dev__set_flag", {"name": n}))'
    const undeclared = await run(runCtx(), code)
    expect(undeclared.text).toMatch(/工具不可用：dev__set_flag/)
    expect(undeclared.isError).toBe(true)

    approvals.length = 0
    flags.length = 0
    answer = { resolution: 'allowed', scope: 'once' }
    const ok = await run(runCtx(), code, ['dev__set_flag', 'dev__get_item'])
    expect(ok.text).toContain('ok:c')
    expect(flags).toEqual(['a', 'b', 'c'])
    expect(approvals).toHaveLength(1)
    expect(approvals[0]).toMatchObject({ tool: 'mcp__gateway__run_script', title: '脚本将调用：dev__set_flag', options: ['allow_once', 'deny'] })

    answer = { resolution: 'denied', reason: '不要改' }
    const denied = await run(runCtx(), code, ['dev__set_flag'])
    expect(denied.text).toMatch(/拒绝了这个脚本：不要改/)
    expect(flags).toHaveLength(3)

    // 自动模式不审批
    approvals.length = 0
    await run(runCtx({ permissionMode: 'auto' }), code, ['dev__set_flag'])
    expect(approvals).toHaveLength(0)

    const wipe = await run(runCtx(), 'print(1)', ['dev__wipe'])
    expect(wipe.text).toMatch(/不能在脚本中调用.*dev__wipe/)
  })

  it('次数上限；脚本出错时返回错误信息；只读模式下工作区只读', async () => {
    const limited = runCtx()
    limited.settings.scripts.maxCalls = 5
    const r = await run(limited, 'from agent_tools import call, ToolError\nn = 0\ntry:\n    for i in range(10):\n        call("dev__get_item", {"i": i}); n += 1\nexcept ToolError as e:\n    print(n, e.message)')
    expect(r.text).toMatch(/5 已达到每个脚本调用工具的上限（5 次）/)

    const bad = await run(runCtx(), 'print("开始")\nraise ValueError("数据不对")')
    expect(bad.isError).toBe(true)
    expect(bad.text).toMatch(/出错退出（退出码 1）/)
    expect(bad.text).toContain('开始')
    expect(bad.text).toContain('ValueError: 数据不对')

    const ro = await run(runCtx({ permissionMode: 'readonly' }), 'from agent_tools import call\nprint(call("dev__get_item", {"i": 2})["value"])\nopen("x.txt", "w").write("x")')
    expect(ro.isError).toBe(true)
    expect(ro.text).toContain('4')
    expect(ro.text).toMatch(/Read-only file system/)
  })
})