| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255 |
- // 后台调用(docs/design/15-background-calls.md):外部工具的长调用转到后台,结束后把结果发给发起方
- import http from 'node:http'
- import type { AddressInfo } from 'node:net'
- 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 { BackgroundCallDto, ProviderStatusDto, TeamDetailDto, TeamEvent, TeamMessage } from '@agent-server/shared'
- import { createApp, type App } from '../src/app'
- import { buildServer } from '../src/http/server'
- import type { AgentProvider, RunContext } from '../src/providers/types'
- type ToolCall = (tool: string, args?: Record<string, unknown>) => Promise<string>
- type Script = (input: string, tool: ToolCall, collab: ToolCall) => Promise<string>
- /** 外部服务的慢工具:等测试放行后才返回 */
- const gates = new Map<string, () => void>()
- const waitGate = (key: string) => new Promise<void>(r => gates.set(key, r))
- const release = async (key: string) => {
- const r = await waitFor(() => gates.get(key))
- gates.delete(key)
- r()
- }
- class ScriptedProvider implements AgentProvider {
- readonly id = 'claude' as const
- scripts: Record<string, Script> = {}
- inputs: { bot: string; text: string; sessionId: string }[] = []
- hang = false
- async run(ctx: RunContext) {
- ctx.setProviderSessionId(`fake-${ctx.sessionId}`)
- const title = app.sessions.getRow(ctx.sessionId).title
- const bot = title.includes(' · ') ? title.split(' · ').at(-1)! : 'session'
- this.inputs.push({ bot, text: ctx.input.text, sessionId: ctx.sessionId })
- if (this.hang) await new Promise((_, reject) => ctx.signal.addEventListener('abort', () => reject(new Error('aborted'))))
- // 与真实引擎一样经 HTTP 连接网关(测试服务的端口与配置不同,替换主机和端口)
- const url = new URL(ctx.gateway.url)
- const client = new Client({ name: 'fake', version: '1' })
- await client.connect(new StreamableHTTPClientTransport(new URL(`${url.pathname}${url.search}`, base), { requestInit: { headers: { Authorization: `Bearer ${ctx.gateway.token}` } } }))
- const tool: ToolCall = async (name, args = {}) => {
- const r = await client.callTool({ name: `slowsvc__${name}`, arguments: args }, undefined, { timeout: 60_000 })
- return (r.content as { text: string }[])[0]!.text
- }
- const gctx = app.gatewayTokens.get(ctx.gateway.token)!
- const collab: ToolCall = async (name, args = {}) => {
- const t = gctx.extraTools!.find(x => x.name === name)!
- const r = await t.handler(t.schema.parse(args) as Record<string, unknown>, gctx)
- return String(r.content[0]?.text ?? '')
- }
- try {
- const reply = await (this.scripts[bot] ?? (async () => 'ok'))(ctx.input.text, tool, collab)
- ctx.emit('message.completed', { messageId: `m${this.inputs.length}`, text: reply })
- } finally {
- await client.close()
- }
- return { usage: { inputTokens: 1, outputTokens: 1, cacheReadTokens: 0, cacheWriteTokens: 0, costUsd: 0.001 } }
- }
- async status(): Promise<ProviderStatusDto> {
- return { id: 'claude', enabled: true, available: true, loggedIn: true, account: 'test', cliPath: '', version: 'x', models: ['m'], defaultModel: 'm', message: null }
- }
- }
- let app: App
- let server: Awaited<ReturnType<typeof buildServer>>
- let base: string
- let upstream: http.Server
- let apiKey: string
- let botKey: string
- let serverId: string
- let adminCookie: Record<string, string>
- const fake = new ScriptedProvider()
- const auth = (key = apiKey) => ({ authorization: `Bearer ${key}` })
- async function waitFor<T>(fn: () => T | undefined | false | null, ms = 6000): Promise<T> {
- const end = Date.now() + ms
- while (Date.now() < end) {
- const v = fn()
- if (v) return v
- await new Promise(r => setTimeout(r, 20))
- }
- throw new Error('timeout')
- }
- function startUpstream(): Promise<string> {
- upstream = http.createServer(async (req, res) => {
- const mcp = new McpServer({ name: 'slow', version: '1.0.0' })
- mcp.registerTool('export', { description: '导出(很慢)', inputSchema: { key: z.string() }, annotations: { readOnlyHint: true } }, async ({ key }) => {
- await waitGate(key)
- return { content: [{ type: 'text', text: key === 'big' ? `结果:${'数'.repeat(9000)}` : `导出完成:${key}` }] }
- })
- mcp.registerTool('quick', { description: '很快', inputSchema: {}, annotations: { readOnlyHint: true } }, async () => ({ content: [{ type: 'text', text: '马上好了' }] }))
- 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`)))
- }
- beforeAll(async () => {
- const log = pino({ level: 'silent' })
- app = createApp(log, { providers: { claude: fake } })
- 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: 'slowsvc', title: '慢服务', description: '', transport: 'http', enabled: true, url, headers: [], command: '', args: [], env: [],
- allowNetwork: false, readPaths: [], writePaths: [],
- }).id
- await app.mcpServers.refresh(serverId)
- const builtin = app.profiles.getBuiltin()
- const { id: _id, version: _v, updatedAt: _u, builtin: _b, builtinKey: _k, ...input } = builtin
- const profile = app.profiles.create({ ...input, name: '带慢服务', isDefault: false, mcpServerIds: [serverId] })
- apiKey = app.apiKeys.create('平台A', { allowedProfileIds: [profile.id], defaultProfileId: profile.id }).key
- for (const [name, kind] of [['组长', 'lead'], ['导出员', 'worker']] as const) app.bots.create({ name, avatar: 'kl:bot:blue', kind, description: name, profileId: profile.id, enabled: true })
- botKey = app.botKeys.create('平台A', app.bots.list().map(b => b.id)).key
- app.settings.update({ backgroundCalls: { autoAfterSec: 1 } })
- const login = await server.inject({ method: 'POST', url: '/admin/v1/auth/login', payload: { password: 'test-password' } })
- adminCookie = { [login.cookies[0]!.name]: login.cookies[0]!.value }
- })
- afterAll(async () => {
- app.backgroundCalls.stop()
- await app.mcpUpstream.closeAll()
- await server.close()
- upstream.close()
- })
- const newSession = async () => (await server.inject({ method: 'POST', url: '/v1/sessions', headers: auth(), payload: {} })).json().id as string
- const calls = async (sid: string) => (await server.inject({ method: 'GET', url: `/v1/sessions/${sid}/background-calls`, headers: auth() })).json() as BackgroundCallDto[]
- const send = (sid: string, text: string) => server.inject({ method: 'POST', url: `/v1/sessions/${sid}/messages`, headers: auth(), payload: { text } })
- const idle = (sid: string) => !app.runs.activeForSession(sid)
- describe('会话中的后台调用', () => {
- it('很快返回的调用照常返回;超过时限的转为后台,完成后把结果作为新消息发到会话', async () => {
- const sid = await newSession()
- let reply = ''
- fake.scripts.session = async (input, tool) => {
- if (input.startsWith('【后台调用')) return '收到结果'
- reply = await tool(input.includes('快') ? 'quick' : 'export', input.includes('快') ? {} : { key: 'q3' })
- return reply
- }
- await send(sid, '快一点')
- await waitFor(() => idle(sid) && reply)
- expect(reply).toBe('马上好了')
- expect(await calls(sid)).toHaveLength(0)
- reply = ''
- await send(sid, '导出 Q3')
- await waitFor(() => idle(sid) && reply)
- expect(reply).toMatch(/已转为后台执行(后台调用 bg_/)
- const [running] = await calls(sid)
- expect(running).toMatchObject({ status: 'running', serverTitle: '慢服务', tool: 'export', args: { key: 'q3' } })
- const before = fake.inputs.length
- await release('q3')
- const delivered = await waitFor(() => fake.inputs.slice(before).find(i => i.sessionId === sid))
- expect(delivered.text).toMatch(/^【后台调用完成】慢服务 · export(bg_.+,耗时 \d+ 分钟)/)
- expect(delivered.text).toContain('结果:\n导出完成:q3')
- await waitFor(() => idle(sid))
- expect((await calls(sid))[0]).toMatchObject({ status: 'completed', result: '导出完成:q3', delivered: true })
- })
- it('会话正在执行时,等这一轮结束再通知;较长的结果保存为文件', async () => {
- const sid = await newSession()
- fake.scripts.session = async (input, tool) => (input.startsWith('【后台调用') || input === '别的事' ? 'ok' : tool('export', { key: 'big' }))
- await send(sid, '导出大文件')
- await waitFor(() => idle(sid) && app.backgroundCalls.list({ sessionId: sid })[0])
- fake.hang = true
- await send(sid, '别的事')
- await waitFor(() => !idle(sid))
- await release('big')
- await waitFor(() => app.backgroundCalls.list({ sessionId: sid })[0]?.status === 'completed')
- const before = fake.inputs.length
- await new Promise(r => setTimeout(r, 100))
- expect(fake.inputs.length).toBe(before)
- expect(app.backgroundCalls.list({ sessionId: sid })[0]!.delivered).toBe(false)
- fake.hang = false
- app.runs.interrupt(app.runs.activeForSession(sid)!.runId, sid)
- const delivered = await waitFor(() => fake.inputs.slice(before).find(i => i.text.startsWith('【后台调用完成】')))
- expect(delivered.text).toMatch(/完整结果已保存到 \.agent\/mcp\/background-export-bg_.+\.txt/)
- await waitFor(() => idle(sid))
- })
- it('标记为总是后台执行的工具立即转到后台;平台可以取消,AI 收到取消的消息', async () => {
- const r = await server.inject({ method: 'PUT', url: `/admin/v1/mcp/servers/${serverId}/tools/export/background`, cookies: adminCookie, payload: { background: true } })
- expect(r.json().tools.find((t: { name: string }) => t.name === 'export').background).toBe(true)
- const sid = await newSession()
- let reply = ''
- const t0 = Date.now()
- fake.scripts.session = async (input, tool) => (input.startsWith('【后台调用') ? 'ok' : (reply = await tool('export', { key: 'never' })))
- await send(sid, '导出')
- await waitFor(() => idle(sid) && reply)
- expect(reply).toContain('已转为后台执行')
- expect(Date.now() - t0).toBeLessThan(900)
- const [call] = await calls(sid)
- // 其他平台不能取消
- const other = app.apiKeys.create('平台B', { allowedProfileIds: [app.profiles.getBuiltin().id], defaultProfileId: app.profiles.getBuiltin().id }).key
- expect((await server.inject({ method: 'POST', url: `/v1/background-calls/${call!.id}/cancel`, headers: auth(other) })).statusCode).toBe(404)
- const before = fake.inputs.length
- const cancelled = await server.inject({ method: 'POST', url: `/v1/background-calls/${call!.id}/cancel`, headers: auth() })
- expect(cancelled.json()).toMatchObject({ status: 'cancelled', error: '成员取消了这次调用' })
- const delivered = await waitFor(() => fake.inputs.slice(before).find(i => i.sessionId === sid))
- expect(delivered.text).toMatch(/^【后台调用已取消】/)
- expect((await server.inject({ method: 'POST', url: `/v1/background-calls/${call!.id}/cancel`, headers: auth() })).statusCode).toBe(409)
- await waitFor(() => idle(sid))
- gates.delete('never')
- await server.inject({ method: 'PUT', url: `/admin/v1/mcp/servers/${serverId}/tools/export/background`, cookies: adminCookie, payload: { background: false } })
- })
- })
- describe('工作组中的后台调用', () => {
- it('组员在子任务中发起的调用转到后台时,子任务等待;结果到达后组员继续并汇报给组长', async () => {
- const [lead, worker] = [app.bots.list().find(b => b.kind === 'lead')!, app.bots.list().find(b => b.kind === 'worker')!]
- const team = (await server.inject({ method: 'POST', url: '/bot/v1/teams', headers: auth(botKey), payload: { human: { id: 'u', name: '张三' }, leadBotId: lead.id, workerBotIds: [worker.id] } })).json() as TeamDetailDto
- fake.scripts['组长'] = async (input, _tool, collab) => {
- if (input.includes('子任务 #') && input.includes('已完成')) {
- await collab('job_complete', { summary: '导出好了' })
- return '完成'
- }
- if (input.includes('导出')) {
- await collab('job_start', { title: '导出', plan: '' })
- await collab('task_create', { bot_id: worker.id, title: '导出数据', instructions: '调用 export' })
- }
- return '安排了'
- }
- fake.scripts['导出员'] = async (input, tool, collab) => {
- if (input.includes('继续完成这个子任务')) {
- expect(input).toContain('【后台调用完成】慢服务 · export')
- await collab('task_complete', { result: input.includes('导出完成:team') ? '已导出' : '?' })
- return '汇报了'
- }
- return tool('export', { key: 'team' })
- }
- await server.inject({ method: 'POST', url: `/bot/v1/teams/${team.id}/messages`, headers: auth(botKey), payload: { text: '帮我导出' } })
- const task = await waitFor(() => app.teams.jobs.tasks(team.id, { status: ['background'] })[0])
- expect(task.title).toBe('导出数据')
- const list = (await server.inject({ method: 'GET', url: `/bot/v1/teams/${team.id}/background-calls`, headers: auth(botKey) })).json() as BackgroundCallDto[]
- expect(list[0]).toMatchObject({ status: 'running', teamId: team.id, botId: worker.id, taskId: task.id })
- // 组长没有被唤醒去处理"已转为后台"的回复
- const messages = () => app.teams.events.list(team.id, 0, 5000, ['message.created']).map(e => (e as TeamEvent & { type: 'message.created' }).payload as TeamMessage)
- expect(messages().some(m => m.card?.kind === 'task' && m.card.event === 'completed')).toBe(false)
- await release('team')
- const job = await waitFor(() => app.teams.jobs.jobs(team.id)[0]?.status === 'completed' && app.teams.jobs.jobs(team.id)[0])
- expect(job.summary).toBe('导出好了')
- expect(app.teams.jobs.tasks(team.id)[0]).toMatchObject({ status: 'completed', result: '已导出', followUp: null })
- const overview = (await server.inject({ method: 'GET', url: '/admin/v1/background-calls', cookies: adminCookie })).json()
- expect(overview.find((x: { id: string }) => x.id === list[0]!.id)).toMatchObject({ status: 'completed', targetTitle: team.title })
- })
- })
|