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

// MCP 支持完善 · 阶段 B：长任务（MCP Tasks）与后台调用合并（docs/design/16-mcp-support.md §3）
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 { InMemoryTaskMessageQueue, InMemoryTaskStore } from '@modelcontextprotocol/sdk/experimental/tasks/stores/in-memory.js'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
import type { BackgroundCallDto, ProviderStatusDto } from '@agent-server/shared'
import { createApp, type App } from '../src/app'
import { buildServer } from '../src/http/server'
import { taskTiming } from '../src/mcp-gateway/tasks'
import type { AgentProvider, RunContext } from '../src/providers/types'

type ToolCall = (tool: string, args?: Record<string, unknown>) => Promise<string>

class ScriptedProvider implements AgentProvider {
  readonly id = 'claude' as const
  script: (input: string, tool: ToolCall) => Promise<string> = async () => 'ok'
  inputs: { text: string; sessionId: string }[] = []
  async run(ctx: RunContext) {
    ctx.setProviderSessionId(`fake-${ctx.sessionId}`)
    this.inputs.push({ text: ctx.input.text, sessionId: ctx.sessionId })
    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: `jobs__${name}`, arguments: args }, undefined, { timeout: 60_000 })
      return (r.content as { text: string }[]).map(c => c.text).join('\n')
    }
    try {
      ctx.emit('message.completed', { messageId: `m${this.inputs.length}`, text: await this.script(ctx.input.text, tool) })
    } 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
const fake = new ScriptedProvider()
const auth = () => ({ authorization: `Bearer ${apiKey}` })

async function waitFor<T>(fn: () => T | undefined | false | null, ms = 8000): 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')
}

/** 外部服务：任务不按 MCP 会话归属（无状态服务），重新连接后也能查询；任务由测试放行后完成 */
const taskStore = new InMemoryTaskStore()
const queue = new InMemoryTaskMessageQueue()
const gates = new Map<string, (r: { text: string; isError?: boolean }) => void>()
const created: string[] = []
function startUpstream(): Promise<string> {
  upstream = http.createServer(async (req, res) => {
    const mcp = new McpServer({ name: 'jobs', version: '1.0.0' }, { capabilities: { tasks: { requests: { tools: { call: {} } }, cancel: {} } }, taskStore, taskMessageQueue: queue })
    mcp.experimental.tasks.registerToolTask('run_job', {
      description: '执行作业（可能很久）', inputSchema: { key: z.string() }, execution: { taskSupport: 'required' },
    }, {
      async createTask({ key }, { taskStore: store, taskRequestedTtl }) {
        const task = await store.createTask({ ttl: taskRequestedTtl, pollInterval: 100 })
        created.push(task.taskId)
        const finish = (r: { text: string; isError?: boolean }) => void taskStore.storeTaskResult(task.taskId, r.isError ? 'failed' : 'completed', { content: [{ type: 'text', text: r.text }], ...(r.isError ? { isError: true } : {}) })
        if (key.startsWith('quick')) setTimeout(() => finish({ text: `作业完成：${key}` }), 200)
        else gates.set(key, finish)
        return { task, _meta: { 'io.modelcontextprotocol/model-immediate-response': `作业 ${key} 已排队，预计 1 小时` } }
      },
      getTask: async (_args, { taskId, taskStore: store }) => (await store.getTask(taskId))!,
      getTaskResult: async (_args, { taskId, taskStore: store }) => await store.getTaskResult(taskId) as never,
    })
    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 finish = async (key: string, r: { text: string; isError?: boolean }) => (await waitFor(() => gates.get(key)))(r)

Object.assign(taskTiming, { minPollMs: 50, defaultPollMs: 100, foregroundPollMs: 50, tickMs: 50 })

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()
  const serverId = app.mcpServers.create({
    name: 'jobs', title: '作业服务', description: '', transport: 'http', enabled: true, url, headers: [], command: '', args: [], env: [],
    allowNetwork: false, readPaths: [], writePaths: [],
  }).id
  await app.mcpServers.refresh(serverId)
  app.mcpServers.setToolPolicy(serverId, 'run_job', 'auto')
  const { id: _id, version: _v, updatedAt: _u, builtin: _b, builtinKey: _k, ...input } = app.profiles.getBuiltin()
  const profile = app.profiles.create({ ...input, name: '带作业服务', isDefault: false, mcpServerIds: [serverId] })
  apiKey = app.apiKeys.create('平台T', { allowedProfileIds: [profile.id], defaultProfileId: profile.id }).key
  app.settings.update({ backgroundCalls: { autoAfterSec: 1 } })
})
afterAll(async () => {
  app.mcpTasks.stop()
  app.backgroundCalls.stop()
  await app.mcpUpstream.closeAll()
  await server.close()
  upstream.closeAllConnections()
  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 () => {
    expect(app.mcpServers.list().find(s => s.name === 'jobs')!.tools[0]).toMatchObject({ name: 'run_job', taskSupport: 'required' })
    const sid = await newSession()
    let reply = ''
    fake.script = async (input, tool) => (reply = await tool('run_job', { key: 'quick1' }))
    await send(sid, '跑一下')
    await waitFor(() => idle(sid) && reply)
    expect(reply).toBe('作业完成：quick1')
    expect(await calls(sid)).toHaveLength(0)
  })

  it('超过时限的任务转为后台（带上服务的说明），按任务 id 查询，完成后把结果发到会话', async () => {
    const sid = await newSession()
    let reply = ''
    fake.script = async (input, tool) => (input.startsWith('【后台调用') ? '收到' : (reply = await tool('run_job', { key: 'slow1' })))
    await send(sid, '跑个大的')
    await waitFor(() => idle(sid) && reply, 10_000)
    expect(reply).toMatch(/已转为后台执行（后台调用 bg_.+）/)
    expect(reply).toContain('服务的说明：作业 slow1 已排队，预计 1 小时')
    const [running] = await calls(sid)
    expect(running).toMatchObject({ status: 'running', mode: 'task', tool: 'run_job' })

    const before = fake.inputs.length
    await finish('slow1', { text: '报表已生成：1024 行' })
    const delivered = await waitFor(() => fake.inputs.slice(before).find(i => i.sessionId === sid))
    expect(delivered.text).toMatch(/^【后台调用完成】作业服务 · run_job/)
    expect(delivered.text).toContain('报表已生成：1024 行')
    await waitFor(() => idle(sid))
    expect((await calls(sid))[0]).toMatchObject({ status: 'completed', delivered: true })
  }, 20_000)

  it('任务失败按失败通知；平台取消时通知服务取消任务', async () => {
    const sid = await newSession()
    fake.script = async (input, tool) => (input.startsWith('【后台调用') ? '收到' : tool('run_job', { key: input }))
    await send(sid, 'fail1')
    await waitFor(() => idle(sid) && app.backgroundCalls.list({ sessionId: sid })[0], 10_000)
    await finish('fail1', { text: '数据源不可用', isError: true })
    await waitFor(() => app.backgroundCalls.list({ sessionId: sid })[0]?.status === 'failed')
    expect(app.backgroundCalls.list({ sessionId: sid })[0]!.error).toContain('数据源不可用')

    await waitFor(() => idle(sid))
    await send(sid, 'cancel1')
    const row = await waitFor(() => idle(sid) && app.backgroundCalls.list({ sessionId: sid }).find(c => c.status === 'running'), 10_000)
    const res = await server.inject({ method: 'POST', url: `/v1/background-calls/${row.id}/cancel`, headers: auth() })
    expect(res.statusCode).toBe(200)
    const taskId = created.at(-1)!
    await expect.poll(async () => (await taskStore.getTask(taskId))?.status).toBe('cancelled')
  }, 25_000)

  it('服务重启后按任务 id 继续查询', async () => {
    const sid = await newSession()
    fake.script = async (input, tool) => (input.startsWith('【后台调用') ? '收到' : tool('run_job', { key: 'restart1' }))
    await send(sid, '跑')
    await waitFor(() => idle(sid) && app.backgroundCalls.list({ sessionId: sid })[0], 10_000)

    // 模拟重启：旧实例停止查询、断开连接，新实例接手同一个数据库
    app.mcpTasks.stop()
    app.backgroundCalls.stop()
    await app.mcpUpstream.closeAll()
    const app2 = createApp(pino({ level: 'silent' }), { db: app.db, providers: { claude: fake } })
    try {
      expect(app2.backgroundCalls.list({ sessionId: sid })[0]).toMatchObject({ status: 'running', mode: 'task' })
      const before = fake.inputs.length
      await finish('restart1', { text: '重启后也拿到了' })
      const delivered = await waitFor(() => fake.inputs.slice(before).find(i => i.sessionId === sid), 10_000)
      expect(delivered.text).toContain('重启后也拿到了')
    } finally {
      app2.mcpTasks.stop()
      app2.backgroundCalls.stop()
      await app2.mcpUpstream.closeAll()
    }
  }, 25_000)
})