diff --git a/.changeset/print-wire-flush-on-exit.md b/.changeset/print-wire-flush-on-exit.md new file mode 100644 index 00000000000..6e649d403b6 --- /dev/null +++ b/.changeset/print-wire-flush-on-exit.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix print mode (`kimi -p`) losing session records when the run exits on an error or a termination signal. diff --git a/apps/kimi-code/src/cli/v2/run-v2-print.ts b/apps/kimi-code/src/cli/v2/run-v2-print.ts index f329c0e06c2..72ed2e9443f 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -22,6 +22,7 @@ import { IAgentCronService, IAgentGoalService, IAgentLifecycleService, + IAgentLoopService, IAgentPermissionModeService, IAgentProfileService, IAgentPromptService, @@ -30,6 +31,7 @@ import { IBootstrapService, IConfigService, IEventBus, + IEventDispatcher, IOAuthToolkit, ISessionIndex, ISessionManager, @@ -122,6 +124,8 @@ import { const PROMPT_UI_MODE = 'print'; /** Re-check `goalActive` at least this often while waiting for goal turns. */ const GOAL_WAIT_POLL_MS = 250; +/** Re-check each agent's prompt queue while waiting for it to drain at exit. */ +const PROMPT_QUIESCE_POLL_MS = 10; /** * Slack on top of a scheduled cron fire time while waiting for the steered * turn: covers the 1s tick poll interval plus fire → inject → turn-launch @@ -196,6 +200,9 @@ export async function runV2Print( } let restorePermission = async (): Promise => {}; + let quiesceAgents = async (): Promise => {}; + let releaseQuiescence: (() => void) | undefined; + let flushWires = async (): Promise => {}; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; let telemetryService: ITelemetryService | undefined; @@ -205,12 +212,25 @@ export async function runV2Print( setCrashPhase('shutdown'); try { await restorePermission(); + // A termination signal can arrive mid-turn: cancel turns and wait for idle agents first. + await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { - if (telemetryService !== undefined) { - await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS); + try { + // Concurrent so the phases' allowances cannot sum past PROMPT_CLEANUP_TIMEOUT_MS. + await Promise.all([ + // The turn's tail records reach the journal only via the wire's + // async persist queue; process.exit must not cut off that queue. + raceWithTimeout(flushWires(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}), + telemetryService !== undefined + ? raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS) + : Promise.resolve(), + shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }).catch(() => {}), + ]); + app.dispose(); + } finally { + // Keep producers frozen until the journals are drained and disposed. + releaseQuiescence?.(); } - await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }).catch(() => {}); - app.dispose(); } })()); await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS); @@ -254,6 +274,10 @@ export async function runV2Print( const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr); restorePermission = resolved.restorePermission; + quiesceAgents = async () => { + releaseQuiescence = await quiesceSessionAgents(resolved.session, resolved.agent); + }; + flushWires = () => flushSessionWires(resolved.session, resolved.agent); telemetryService.setContext({ session_id: resolved.session.id, model: resolved.telemetryModel }); setTelemetryContext({ sessionId: resolved.session.id }); @@ -869,6 +893,117 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number { return count; } +/** Every agent handle in the session; the main agent is included explicitly since the lifecycle list skips `closing` agents. */ +function collectSessionAgentHandles( + session: ISessionScopeHandle, + mainAgent: IAgentScopeHandle, +): IAgentScopeHandle[] { + const agentManager = session.accessor.get(IAgentLifecycleService); + const handles = new Set([mainAgent]); + for (const agent of agentManager.list()) { + const handle = agentManager.handleOf(agent.agentId); + if (handle !== undefined) handles.add(handle); + } + return [...handles]; +} + +/** + * Stop producers, drain prompts, and cancel turns so closing records exist + * before the wire flush; returns a release holding a guard per loop. + */ +async function quiesceSessionAgents( + session: ISessionScopeHandle, + mainAgent: IAgentScopeHandle, +): Promise<(() => void) | undefined> { + const handles = collectSessionAgentHandles(session, mainAgent); + const promptServices = handles.flatMap((handle) => { + try { + return [handle.accessor.get(IAgentPromptService)]; + } catch { + // A torn-down agent scope has no prompt service to drain or observe. + return []; + } + }); + const loops = handles.flatMap((handle) => { + try { + return [handle.accessor.get(IAgentLoopService)]; + } catch { + // A torn-down agent scope has no loop to quiesce. + return []; + } + }); + // Task producers bypass the prompt queue and dispatch termination records + // straight to the wire; stop them first so the flush can persist those. + await Promise.allSettled( + handles.flatMap((handle) => { + try { + return [handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed')]; + } catch { + return []; + } + }), + ); + // Repeat until every queue is empty and every loop freezable: a prompt can + // still surface from the launch window or a cancelled turn's settle chain. + for (;;) { + await Promise.allSettled(promptServices.map((service) => service.drain())); + for (const loop of loops) { + for (const turnId of loop.status().pendingTurnIds) loop.cancel(turnId); + loop.cancel(); + } + await Promise.allSettled(loops.map((loop) => loop.settled())); + const guards: { dispose(): void }[] = []; + let frozen = true; + for (const loop of loops) { + let guard: { dispose(): void } | undefined; + try { + guard = loop.tryAcquireQuiescence(); + } catch { + // A disposed loop cannot accept new submissions; it needs no guard. + continue; + } + if (guard === undefined) { + frozen = false; + break; + } + guards.push(guard); + } + const busy = promptServices.some((service) => { + try { + const snapshot = service.list(); + return ( + snapshot.launching || + snapshot.active !== undefined || + snapshot.pending.length > 0 + ); + } catch { + return false; + } + }); + if (frozen && !busy) { + return () => { + for (const guard of guards) guard.dispose(); + }; + } + for (const guard of guards) guard.dispose(); + await new Promise((resolve) => { + setTimeout(resolve, PROMPT_QUIESCE_POLL_MS); + }); + } +} + +/** Flush every session agent's wire journal; each flush settles independently. */ +async function flushSessionWires( + session: ISessionScopeHandle, + mainAgent: IAgentScopeHandle, +): Promise { + await Promise.allSettled( + collectSessionAgentHandles(session, mainAgent).map((handle) => + handle.accessor.get(IEventDispatcher).flush(), + ), + ); +} + async function drainBackgroundTasks( session: ISessionScopeHandle, ceilingS: number | undefined, diff --git a/apps/kimi-code/test/cli/v2-run-print.test.ts b/apps/kimi-code/test/cli/v2-run-print.test.ts index fdb114de2aa..e3b12c1b704 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -8,6 +8,7 @@ import { IAgentCronService, IAgentGoalService, IAgentLifecycleService, + IAgentLoopService, IAgentPermissionModeService, IAgentProfileService, IAgentPromptService, @@ -17,6 +18,7 @@ import { IBootstrapService, IConfigService, IEventBus, + IEventDispatcher, IFileSystemStorageService, IOAuthToolkit, ISessionIndex, @@ -181,11 +183,23 @@ function makeFakeHarness() { }), }; }), + drain: vi.fn(async () => {}), + list: vi.fn(() => ({ launching: false, active: undefined, pending: [] })), }, ], - [IAgentTaskService, { list: vi.fn(() => []) }], + [IAgentTaskService, { list: vi.fn(() => []), stopAllOnExit: vi.fn(async () => []) }], [IAgentCronService, { getNextFireTime: vi.fn(() => null) }], [IAgentGoalService, goal], + [IEventDispatcher, { flush: vi.fn(async () => {}) }], + [ + IAgentLoopService, + { + status: vi.fn(() => ({ state: 'idle', pendingTurnIds: [] })), + cancel: vi.fn(() => false), + settled: vi.fn(async () => {}), + tryAcquireQuiescence: vi.fn(() => ({ dispose: vi.fn() })), + }, + ], [ IAgentScopeContext, makeAgentScopeContext({ agentId: 'main', agentScope: 'agents/main' }), @@ -274,7 +288,7 @@ function makeFakeHarness() { ], ]); const app = fakeScope('app', appServices); - return { app, agent, session, agentServices, appServices, profileState }; + return { app, agent, session, agentServices, sessionServices, appServices, profileState }; } describe('runV2Print', () => { @@ -613,4 +627,200 @@ describe('runV2Print', () => { expect(initOrder).toBeDefined(); expect(reconcileOrder).toBeGreaterThan(initOrder!); }); + + it('flushes the wire journal before disposing the app', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agentServices } = makeFakeHarness(); + + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); + + await runV2Print(opts() as never, '1.2.3-test', { stdout, stderr }); + + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + expect(dispatcher.flush).toHaveBeenCalled(); + const flushOrder = dispatcher.flush.mock.invocationCallOrder[0]; + const disposeOrder = app.dispose.mock.invocationCallOrder[0]; + expect(flushOrder).toBeDefined(); + expect(disposeOrder).toBeGreaterThan(flushOrder!); + }); + + it('flushes the wire journal when the turn fails', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agentServices } = makeFakeHarness(); + + const promptService = agentServices.get(IAgentPromptService) as { + enqueue: ReturnType; + }; + promptService.enqueue.mockResolvedValueOnce({ + launched: Promise.resolve({ + id: 1, + result: Promise.resolve({ + type: 'failed', + error: { code: 'provider.overloaded', message: 'llm request failed' }, + }), + }), + }); + + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); + + await expect(runV2Print(opts() as never, '1.2.3-test', { stdout, stderr })).rejects.toThrow( + 'provider.overloaded: llm request failed', + ); + + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + expect(dispatcher.flush).toHaveBeenCalled(); + const flushOrder = dispatcher.flush.mock.invocationCallOrder[0]; + const disposeOrder = app.dispose.mock.invocationCallOrder[0]; + expect(disposeOrder).toBeGreaterThan(flushOrder!); + }); + + it('does not let a wire flush failure mask the turn outcome', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agentServices } = makeFakeHarness(); + + const promptService = agentServices.get(IAgentPromptService) as { + enqueue: ReturnType; + }; + promptService.enqueue.mockResolvedValueOnce({ + launched: Promise.resolve({ + id: 1, + result: Promise.resolve({ + type: 'failed', + error: { code: 'provider.overloaded', message: 'llm request failed' }, + }), + }), + }); + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + dispatcher.flush.mockRejectedValueOnce(new Error('disk full')); + + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); + + await expect(runV2Print(opts() as never, '1.2.3-test', { stdout, stderr })).rejects.toThrow( + 'provider.overloaded: llm request failed', + ); + expect(app.dispose).toHaveBeenCalled(); + }); + + it('cancels and settles the active turn before flushing on a termination signal', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agentServices } = makeFakeHarness(); + + const order: string[] = []; + const loop = agentServices.get(IAgentLoopService) as { + status: ReturnType; + cancel: ReturnType; + settled: ReturnType; + tryAcquireQuiescence: ReturnType; + }; + loop.status.mockReturnValue({ state: 'running', pendingTurnIds: [] }); + loop.cancel.mockImplementation(() => { + if (!order.includes('cancel')) order.push('cancel'); + return true; + }); + loop.settled = vi.fn(async () => { + if (!order.includes('settled')) order.push('settled'); + }); + const guardDispose = vi.fn(); + loop.tryAcquireQuiescence = vi.fn(() => ({ dispose: guardDispose })); + const taskService = agentServices.get(IAgentTaskService) as { + stopAllOnExit: ReturnType; + }; + taskService.stopAllOnExit = vi.fn(async () => { + if (!order.includes('stop')) order.push('stop'); + return []; + }); + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + dispatcher.flush = vi.fn(async () => { + order.push('flush'); + }); + + // A turn still in flight when the signal arrives: the prompt queue reports + // the launch window, then the running prompt, then goes empty. + const promptService = agentServices.get(IAgentPromptService) as { + enqueue: ReturnType; + drain: ReturnType; + list: ReturnType; + }; + let settleTurn!: (result: unknown) => void; + promptService.enqueue.mockResolvedValueOnce({ + launched: Promise.resolve({ + id: 1, + result: new Promise((resolve) => { + settleTurn = resolve; + }), + }), + }); + let promptPhase: 'launching' | 'active' | 'empty' = 'launching'; + promptService.list = vi.fn(() => { + if (promptPhase === 'launching') { + return { launching: true, active: undefined, pending: [] }; + } + if (promptPhase === 'active') { + return { launching: false, active: { id: 'p1' }, pending: [] }; + } + return { launching: false, active: undefined, pending: [] }; + }); + + const handlers = new Map Promise>(); + const fakeProcess = { + once: (signal: string, handler: () => Promise) => { + handlers.set(signal, handler); + }, + off: () => {}, + exit: vi.fn((code?: number) => { + order.push(`exit:${code}`); + }), + }; + + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); + + const run = runV2Print(opts() as never, '1.2.3-test', { + stdout, + stderr, + process: fakeProcess as never, + }); + const outcome = run.catch((error: unknown) => error); + for (let i = 0; i < 100 && !handlers.has('SIGINT'); i++) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + const onSigint = handlers.get('SIGINT')!; + settleTurn({ type: 'cancelled', steps: 0, reason: new Error('aborted') }); + const sigintRun = onSigint(); + // The flush must wait for the prompt queue to empty, even with idle loops. + for (let i = 0; i < 100 && !order.includes('settled'); i++) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(promptService.drain).toHaveBeenCalled(); + expect(order).toEqual(['stop', 'cancel', 'settled']); + promptPhase = 'active'; + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(order).toEqual(['stop', 'cancel', 'settled']); + promptPhase = 'empty'; + await sigintRun; + + expect(order).toEqual(['stop', 'cancel', 'settled', 'flush', 'exit:130']); + // The guard taken during quiesce is only released after app.dispose(). + expect(loop.tryAcquireQuiescence).toHaveBeenCalled(); + const lastGuardRelease = guardDispose.mock.invocationCallOrder.at(-1); + const appDisposeOrder = app.dispose.mock.invocationCallOrder[0]; + expect(lastGuardRelease).toBeGreaterThan(appDisposeOrder!); + expect(await outcome).toBeInstanceOf(Error); + }); }); diff --git a/packages/agent-core-v2/src/agent/prompt/prompt.ts b/packages/agent-core-v2/src/agent/prompt/prompt.ts index 4eb7136d2a1..fc17432448d 100644 --- a/packages/agent-core-v2/src/agent/prompt/prompt.ts +++ b/packages/agent-core-v2/src/agent/prompt/prompt.ts @@ -47,6 +47,7 @@ export interface PromptHandle extends PromptSnapshot { export interface PromptQueueSnapshot { readonly active: PromptSnapshot | undefined; readonly pending: readonly PromptSnapshot[]; + readonly launching: boolean; } export interface PromptPayload { diff --git a/packages/agent-core-v2/src/agent/prompt/promptService.ts b/packages/agent-core-v2/src/agent/prompt/promptService.ts index 81dcac6947e..90a2c421932 100644 --- a/packages/agent-core-v2/src/agent/prompt/promptService.ts +++ b/packages/agent-core-v2/src/agent/prompt/promptService.ts @@ -388,7 +388,7 @@ export class AgentPromptService implements IAgentPromptService { } list(): PromptQueueSnapshot { - return { active: this.active === undefined ? undefined : snapshot(this.active), pending: this.pending.map(snapshot) }; + return { active: this.active === undefined ? undefined : snapshot(this.active), pending: this.pending.map(snapshot), launching: this.launching }; } async steer(promptIds: readonly string[]): Promise { diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index b823e457f8f..c036ba8c07d 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -800,10 +800,17 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { async stopAllOnExit(reason: string): Promise { if (this.keepAliveOnExit()) return []; const active = this.list(true); - await Promise.all( + await Promise.allSettled( active .filter((task) => task.detached === true) - .map((task) => this.suppressTerminalNotification(task.taskId)), + .map((task) => + this.suppressTerminalNotification(task.taskId).catch((error: unknown) => { + this.log.error('terminal notification suppression failed', { + taskId: task.taskId, + error, + }); + }), + ), ); return this.stopAll(reason); } diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index 6afd6746215..635a4ecdff3 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -4,6 +4,7 @@ import { IInstantiationService } from '#/_base/di/instantiation'; import type { InstantiationService } from '#/_base/di/instantiationService'; import { Disposable } from '#/_base/di/lifecycle'; import { Emitter } from '#/_base/event'; +import { onUnexpectedError } from '#/_base/errors/unexpectedError'; import { Error2, ErrorCodes } from '#/errors'; import { LifecycleScope } from '#/app/scopes'; import { @@ -58,6 +59,9 @@ import { let nextAgentId = 0; +const REMOVE_PROMPT_QUIESCE_TIMEOUT_MS = 3_000; +const REMOVE_PROMPT_QUIESCE_POLL_MS = 10; + export class AgentLifecycleService extends Disposable implements IAgentLifecycleService { declare readonly _serviceBrand: undefined; private readonly roster = new Map(); @@ -365,16 +369,46 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle const compactionSettled = compaction?.promise.catch(() => undefined) ?? Promise.resolve(); const reason = abortError('Agent removed'); const prompt = handle.accessor.get(IAgentPromptService); - for (const turnId of loop.status().pendingTurnIds) { - loop.cancel(turnId, reason); - } - loop.cancel(undefined, reason); if (compaction !== null && !compaction.abortController.signal.aborted) { compaction.abortController.abort(reason); } - await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]); - managed.killSpace(); - await handle.dispose(); + const promptIdleDeadline = Date.now() + REMOVE_PROMPT_QUIESCE_TIMEOUT_MS; + let releaseQuiescence: (() => void) | undefined; + for (;;) { + for (const turnId of loop.status().pendingTurnIds) { + loop.cancel(turnId, reason); + } + loop.cancel(undefined, reason); + await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]); + let idle = true; + try { + const snapshot = prompt.list(); + idle = + !snapshot.launching && snapshot.active === undefined && snapshot.pending.length === 0; + } catch { + idle = true; + } + if (idle) { + try { + const guard = loop.tryAcquireQuiescence(); + if (guard !== undefined) { + releaseQuiescence = () => guard.dispose(); + break; + } + } catch { + break; + } + } + if (Date.now() >= promptIdleDeadline) break; + await new Promise((resolve) => setTimeout(resolve, REMOVE_PROMPT_QUIESCE_POLL_MS)); + } + try { + await handle.accessor.get(IEventDispatcher).flush().catch(onUnexpectedError); + managed.killSpace(); + await handle.dispose(); + } finally { + releaseQuiescence?.(); + } if (this.roster.get(agent.agentId) === managed) this.roster.delete(agent.agentId); this.onDidCloseEmitter.fire(agent); } diff --git a/packages/agent-core-v2/test/agent/prompt/promptService.test.ts b/packages/agent-core-v2/test/agent/prompt/promptService.test.ts index c1eece1b3f9..3733283b4b0 100644 --- a/packages/agent-core-v2/test/agent/prompt/promptService.test.ts +++ b/packages/agent-core-v2/test/agent/prompt/promptService.test.ts @@ -245,7 +245,7 @@ describe('AgentPromptService', () => { it('keeps injections outside the prompt queue', async () => { const { prompt } = harness(); await prompt.inject({ ...message('system'), origin: { kind: 'injection', variant: 'test' } }); - expect(prompt.list()).toEqual({ active: undefined, pending: [] }); + expect(prompt.list()).toEqual({ active: undefined, pending: [], launching: false }); }); it('settles blocked prompts', async () => { @@ -258,6 +258,26 @@ describe('AgentPromptService', () => { expect(completed.map((event) => [event.promptId, event.reason])).toEqual([[handle.id, 'blocked']]); }); + it('marks the launch window as busy in the queue snapshot', async () => { + const { prompt } = harness(); + let releaseHook!: () => void; + prompt.hooks.onBeforeSubmitPrompt.register('gate', async (_ctx, next) => { + await new Promise((resolve) => { + releaseHook = resolve; + }); + await next(); + }); + const enqueued = prompt.enqueue({ message: message('launching') }); + await vi.waitFor(() => { + expect(prompt.list().launching).toBe(true); + }); + expect(prompt.list().active).toBeUndefined(); + expect(prompt.list().pending).toEqual([]); + releaseHook(); + await enqueued; + expect(prompt.list().launching).toBe(false); + }); + it('delivers a blocked prompt’s compression captions right after their host message', async () => { const { prompt, context } = harness(); prompt.hooks.onBeforeSubmitPrompt.register('block', async (ctx, next) => { ctx.block = true; await next(); }); @@ -294,7 +314,7 @@ describe('AgentPromptService', () => { expect(handle.state).toBe('failed'); await expect(handle.launched).resolves.toBeUndefined(); await expect(handle.completion).resolves.toMatchObject({ state: 'failed', result: undefined }); - expect(prompt.list()).toEqual({ active: undefined, pending: [] }); + expect(prompt.list()).toEqual({ active: undefined, pending: [], launching: false }); }); it('replaces an unsupported prompt image with a text notice at the history funnel', async () => { diff --git a/packages/agent-core-v2/test/agent/task/taskService.test.ts b/packages/agent-core-v2/test/agent/task/taskService.test.ts index 9f0a10508ba..c8f6e90bc55 100644 --- a/packages/agent-core-v2/test/agent/task/taskService.test.ts +++ b/packages/agent-core-v2/test/agent/task/taskService.test.ts @@ -544,6 +544,28 @@ describe('AgentTaskService', () => { }); }); + it('stopAllOnExit still stops tasks when suppression persistence fails', async () => { + let writes = 0; + ix.stub(IAtomicDocumentStore, { + get: async () => undefined, + set: async () => { + writes += 1; + if (writes === 1) throw new Error('disk full'); + }, + delete: async () => {}, + list: async () => [], + }); + const svc = ix.get(IAgentTaskService); + const first = svc.registerTask(fakeProcessTask()); + const second = svc.registerTask(fakeProcessTask()); + + const stopped = await svc.stopAllOnExit('Session closed'); + + expect(stopped.map((info) => info.taskId).toSorted()).toEqual([first, second].toSorted()); + expect(svc.getTask(first)?.status).toBe('killed'); + expect(svc.getTask(second)?.status).toBe('killed'); + }); + it('stopAllOnExit leaves tasks running when keepAliveOnExit is set', async () => { stubTaskConfig({ keepAliveOnExit: true }); const svc = ix.get(IAgentTaskService); diff --git a/packages/agent-core-v2/test/agent/undo/undo.test.ts b/packages/agent-core-v2/test/agent/undo/undo.test.ts index d88b161d424..925b16d39cc 100644 --- a/packages/agent-core-v2/test/agent/undo/undo.test.ts +++ b/packages/agent-core-v2/test/agent/undo/undo.test.ts @@ -536,6 +536,7 @@ describe('AgentConversationUndoService', () => { ctx.appendTurnExchange('u2', 'a2'); const list = vi.spyOn(ctx.get(IAgentPromptService), 'list').mockReturnValue({ active: undefined, + launching: false, pending: [ { id: 'queued', diff --git a/packages/agent-core-v2/test/app/gateway/gateway.test.ts b/packages/agent-core-v2/test/app/gateway/gateway.test.ts index c015d127d86..01d1ecbc504 100644 --- a/packages/agent-core-v2/test/app/gateway/gateway.test.ts +++ b/packages/agent-core-v2/test/app/gateway/gateway.test.ts @@ -59,7 +59,7 @@ describe('RestGateway', () => { submit: () => Promise.resolve(undefined), submitSteer: () => Promise.resolve(undefined), steer: () => Promise.resolve([]), - list: () => ({ active: undefined, pending: [] }), + list: () => ({ active: undefined, pending: [], launching: false }), abort: () => true, drain: () => Promise.resolve(), inject: () => Promise.resolve(undefined), diff --git a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index f8895d79ce0..b1c6c05a38b 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -348,11 +348,13 @@ describe('AgentLifecycleService', () => { }), cancel: loopCancel, settled: loopSettled, + tryAcquireQuiescence: vi.fn(() => ({ dispose: vi.fn() })), } as unknown as IAgentLoopService); promptDrain = vi.fn(async () => {}); ix.stub(IAgentPromptService, { _serviceBrand: undefined, drain: promptDrain, + list: () => ({ launching: false, active: undefined, pending: [] }), } as unknown as IAgentPromptService); ix.stub(ITelemetryService, { _serviceBrand: undefined, @@ -509,6 +511,15 @@ describe('AgentLifecycleService', () => { expect(svc.handleOf('main')).toBeUndefined(); }); + it('remove flushes the agent wire journal before disposal', async () => { + const svc = ix.get(IAgentLifecycleService); + await svc.create({ agentId: 'main' }); + const dispatcher = svc.handleOf('main')!.accessor.get(IEventDispatcher); + const flush = vi.spyOn(dispatcher, 'flush'); + await svc.remove(svc.get('main')!); + expect(flush).toHaveBeenCalled(); + }); + it('remove keeps the lifecycle context active through async scope teardown', async () => { const svc = ix.get(IAgentLifecycleService); const bus = ix.get(ISessionEventBus); diff --git a/packages/agent-core-v2/test/session/sessionTitle/agentTitlePromptSourceService.test.ts b/packages/agent-core-v2/test/session/sessionTitle/agentTitlePromptSourceService.test.ts index 0c5bd74722e..182ebb3e73e 100644 --- a/packages/agent-core-v2/test/session/sessionTitle/agentTitlePromptSourceService.test.ts +++ b/packages/agent-core-v2/test/session/sessionTitle/agentTitlePromptSourceService.test.ts @@ -41,7 +41,7 @@ describe('AgentTitlePromptSource', () => { beforeEach(() => { liveMessages = []; - queue = { active: undefined, pending: [] }; + queue = { active: undefined, pending: [], launching: false }; disposables = new DisposableStore(); ix = createServices(disposables, { additionalServices: (reg) => { @@ -60,6 +60,7 @@ describe('AgentTitlePromptSource', () => { liveMessages = [userMessage('one', '第一条')]; queue = { active: undefined, + launching: false, pending: [ { id: 'two', @@ -122,6 +123,7 @@ describe('AgentTitlePromptSource', () => { it('counts a queued prompt already appended to the context only once', async () => { liveMessages = [userMessage('one', '同一条')]; queue = { + launching: false, active: { id: 'one', userMessageId: 'one', @@ -171,6 +173,7 @@ describe('AgentTitlePromptSource', () => { userMessage('two', '进行中的问题'), ]; queue = { + launching: false, active: { id: 'two', userMessageId: 'two',