yaoyao/tests/server/workspaceReconnect.test.ts

// @vitest-environment node
import {afterEach, beforeEach, expect, it, vi} from 'vitest'
import {mkdtempSync, rmSync} from 'node:fs'
import {tmpdir} from 'node:os'
import {join} from 'node:path'
import {randomUUID} from 'node:crypto'
import {WebSocketServer, type WebSocket} from 'ws'
import {WorkspaceStore} from '../../src/server/workspaceStore'
import {WorkspaceRuntime} from '../../src/server/workspaceRuntime'
import {type GatewayTarget, type WorkspaceNodes} from '../../src/server/workspaceGateway'
import {type WorkspacePlugins} from '../../src/server/botPlugins/workspacePlugins'
import {UploadStore} from '../../src/server/uploads'

const owner = 'reconnect-user'
let home: string, store: WorkspaceStore, uploads: UploadStore, runtime: WorkspaceRuntime
let server: WebSocketServer, socket: WebSocket, target: GatewayTarget, nodes: WorkspaceNodes
let running: boolean, authorized: boolean, acknowledgePrompt: boolean, holdCredential: (() => Promise<void>) | undefined
let holdBind: (() => Promise<void>) | undefined
let resumedRuntimeId: string, resumedStoredId: string, historyPaths: string[]
let rpc: Array<{method: string; params: any}>, binds: any[], unbinds: any[], history: any[]
let pluginDispose: ReturnType<typeof vi.fn>, pluginCall: ReturnType<typeof vi.fn>
beforeEach(async () => {
  home = mkdtempSync(join(tmpdir(), 'yaoyao-workspace-reconnect-'))
  store = new WorkspaceStore(home); uploads = new UploadStore(home)
  running = true; authorized = true; acknowledgePrompt = true; holdCredential = undefined; holdBind = undefined
  resumedRuntimeId = 'runtime'; resumedStoredId = 'stored'; historyPaths = []
  rpc = []; binds = []; unbinds = []; history = []
  server = new WebSocketServer({port: 0, host: '127.0.0.1'})
  await new Promise<void>(resolve => server.once('listening', resolve))
  target = {
    url: new URL(`http://127.0.0.1:${(server.address() as {port: number}).port}`), client: {directAgent: undefined},
    session: {
      webSocketCredential: async () => { await holdCredential?.(); return {name: 'ticket', value: 'fixture'} },
      request: async (path: string, options: any) => {
        if (path.endsWith('/bind')) { binds.push(options.body); await holdBind?.() }
        if (path.endsWith('/unbind')) unbinds.push(options.body)
        if (path.startsWith('/api/sessions/')) historyPaths.push(path)
        const value = path.endsWith('/capabilities') ? {ready: true, native_tools: true, in_process: true}
          : path.startsWith('/api/plugins/yaoyao-bot-bridge/') ? {ok: true, native_tools: true}
          : path === '/api/config' ? {terminal: {}} : {messages: history}
        return {status: 200, body: Buffer.from(JSON.stringify(value)), headers: new Headers()}
      },
    },
  } as unknown as GatewayTarget
  nodes = {requireSource: () => {}, target: () => target} as unknown as WorkspaceNodes
  runtime = new WorkspaceRuntime(store, nodes, uploads, () => authorized)
  pluginDispose = vi.fn(async () => {})
  pluginCall = vi.fn(async () => ({message: '插件仍可使用'}))
  runtime.plugins = {selected: () => true, open: async () => ({services: () => [], warnings: () => [],
    catalog: () => [{id: 'plugin_fixture', name: 'plugin_fixture', description: '回归测试', inputSchema: {type: 'object'}}],
    call: pluginCall, dispose: pluginDispose,
  })} as unknown as WorkspacePlugins
  server.on('connection', connection => {
    socket = connection
    connection.send(JSON.stringify({method: 'event', params: {type: 'gateway.ready'}}))
    connection.on('message', raw => {
      const frame = JSON.parse(String(raw)); rpc.push(frame)
      const result = frame.method === 'session.create' ? {session_id: 'runtime', stored_session_id: 'stored', running: false, info: {profile_name: 'default'}}
        : frame.method === 'session.resume' ? {session_id: resumedRuntimeId, stored_session_id: resumedStoredId, running, info: {profile_name: 'default'}}
        : {status: 'streaming'}
      if (frame.method !== 'prompt.submit' || acknowledgePrompt) connection.send(JSON.stringify({id: frame.id, result}))
    })
  })
})
afterEach(async () => {
  runtime.close()
  await new Promise(resolve => setTimeout(resolve, 20))
  for (const client of server.clients) client.terminate()
  await new Promise<void>(resolve => server.close(() => resolve()))
  uploads.close(); store.close(); rmSync(home, {recursive: true, force: true})
})
async function start() {
  const agent = store.createAgent(owner, {name: '恢复测试', profile: 'default'})
  const conversation = store.list<any>(owner, 'conversation').find(c => c.kind === 'direct' && c.memberIds[0] === agent.id)
  const run = runtime.send(owner, conversation.id, {requestId: randomUUID(), content: '只执行一次'})
  await expect.poll(() => rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
  return {run, conversation, agent}
}
async function callback(path: string, body: unknown) {
  const lease = binds[0]
  return fetch(lease.bridge_url + path, {method: 'POST', headers: {Authorization: `Bearer ${lease.token}`, 'Content-Type': 'application/json'}, body: JSON.stringify(body)})
}
function complete(text = '恢复完成') { socket.send(JSON.stringify({method: 'event', params: {type: 'message.complete', session_id: 'runtime', payload: {text}}})) }

it.each([true, false])('keeps the original tool and plugin grants through running recovery (prompt receipt=%s)', async receipt => {
  acknowledgePrompt = receipt
  const {run, conversation, agent} = await start()
  const first = binds[0]
  socket.terminate()
  await expect.poll(() => binds.length).toBe(2)
  expect(binds[1]).toMatchObject({generation: first.generation, token: first.token, bridge_url: first.bridge_url})
  expect(unbinds).toHaveLength(0)
  expect(pluginDispose).not.toHaveBeenCalled()
  const catalog = await (await callback('/tools/list', {})).json()
  const tool = catalog.tools.find((t: any) => t.name === 'workspace_list_agents')
  const result = await (await callback('/tools/call', {toolId: tool.id, arguments: {}, callId: 'after-reconnect'})).json()
  expect(result.structuredContent.agents).toEqual(expect.arrayContaining([expect.objectContaining({id: agent.id})]))
  const plugin = catalog.tools.find((t: any) => t.name === 'plugin_fixture')
  await callback('/tools/call', {toolId: plugin.id, arguments: {}, callId: 'plugin-after-reconnect'})
  expect(pluginCall).toHaveBeenCalledOnce()
  complete()
  await expect.poll(() => store.require<any>(owner, 'run', run.id).status).toBe('complete')
  expect(store.messages(owner, conversation.id).at(-1)?.content).toBe('恢复完成')
  expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
  await expect.poll(() => unbinds.length).toBe(1)
  expect(pluginDispose).toHaveBeenCalledOnce()
})

it('reconciles completion during a disconnect without renewing or resubmitting the finished turn', async () => {
  const {run, conversation} = await start()
  history = [{role: 'user', content: rpc.find(r => r.method === 'prompt.submit')!.params.text}, {role: 'assistant', content: '断线期间完成'}]
  running = false; socket.terminate()
  await expect.poll(() => store.require<any>(owner, 'run', run.id).status).toBe('complete')
  expect(store.messages(owner, conversation.id).at(-1)?.content).toBe('断线期间完成')
  expect(binds).toHaveLength(1)
  expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
})

it('continues reconnecting if the replacement socket closes while the original grant is renewing', async () => {
  const {run} = await start()
  let release!: () => void
  holdBind = () => new Promise<void>(resolve => { release = resolve })
  socket.terminate()
  await expect.poll(() => binds.length).toBe(2)
  socket.terminate()
  await expect.poll(() => server.clients.size).toBe(0)
  holdBind = undefined; release()
  await expect.poll(() => binds.length, {timeout: 3000}).toBe(3)
  expect(new Set(binds.map(b => b.generation)).size).toBe(1)
  complete()
  await expect.poll(() => store.require<any>(owner, 'run', run.id).status).toBe('complete')
  expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
})

it('follows resumed runtime and compacted stored IDs for tool renewal and subsequent history reconciliation', async () => {
  const {run} = await start()
  resumedRuntimeId = 'runtime-after-resume'; resumedStoredId = 'compacted-tip'
  socket.terminate()
  await expect.poll(() => binds.length).toBe(2)
  expect(binds[1]).toMatchObject({session_id: resumedRuntimeId, stored_session_id: resumedStoredId, generation: binds[0].generation})
  expect(store.list<any>(owner, 'binding')[0]).toMatchObject({runtimeId: resumedRuntimeId, storedId: resumedStoredId, aliases: ['stored', 'compacted-tip']})
  running = false
  history = [{role: 'user', content: rpc.find(r => r.method === 'prompt.submit')!.params.text}, {role: 'assistant', content: '压缩后完成'}]
  socket.terminate()
  await expect.poll(() => store.require<any>(owner, 'run', run.id).status).toBe('complete')
  expect(historyPaths).toEqual(['/api/sessions/compacted-tip/messages'])
  expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
})

it.each(['stop', 'revoke', 'close'] as const)('revokes callbacks and cannot revive a detached turn after %s', async action => {
  const {run} = await start()
  socket.terminate()
  await expect.poll(() => store.require<any>(owner, 'run', run.id).error).toContain('正在恢复原执行')
  if (action === 'stop') await runtime.stop(owner, run.id)
  else if (action === 'revoke') authorized = false
  else runtime.close()
  await expect.poll(() => unbinds.length).toBe(1)
  await expect(callback('/tools/list', {})).rejects.toThrow()
  expect(pluginDispose).toHaveBeenCalledOnce()
  expect(binds).toHaveLength(1)
  expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
  if (action === 'stop') expect(store.require<any>(owner, 'run', run.id).status).toBe('interrupted')
})

it('does not create a late WebSocket when the service closes during reconnect credentials', async () => {
  const {run} = await start()
  let release!: () => void
  const credential = new Promise<void>(resolve => { release = resolve })
  let waiting = false
  holdCredential = () => { waiting = true; return credential }
  socket.terminate()
  await expect.poll(() => waiting).toBe(true)
  runtime.close(); release()
  await expect.poll(() => unbinds.length).toBe(1)
  expect(server.clients.size).toBe(0)
  expect(rpc.filter(r => r.method === 'session.resume')).toHaveLength(0)
  expect(store.require<any>(owner, 'run', run.id).status).toBe('uncertain')
})

it('keeps a running turn uncertain after process restart when its original tools cannot be restored', async () => {
  const {run} = await start()
  runtime.close()
  runtime = new WorkspaceRuntime(store, nodes, uploads, () => authorized)
  runtime.start()
  await expect.poll(() => store.require<any>(owner, 'run', run.id).error).toContain('本轮工具授权已丢失')
  expect(store.require<any>(owner, 'run', run.id).status).toBe('uncertain')
  expect(binds).toHaveLength(1)
  expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1)
})