// @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) | undefined let holdBind: (() => Promise) | undefined let resumedRuntimeId: string, resumedStoredId: string, historyPaths: string[] let rpc: Array<{method: string; params: any}>, binds: any[], unbinds: any[], history: any[] let pluginDispose: ReturnType, pluginCall: ReturnType 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(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(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(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(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(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(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(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(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(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(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(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(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(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(owner, 'run', run.id).error).toContain('本轮工具授权已丢失') expect(store.require(owner, 'run', run.id).status).toBe('uncertain') expect(binds).toHaveLength(1) expect(rpc.filter(r => r.method === 'prompt.submit')).toHaveLength(1) })