From 3511b89a48af4f41eb4155d7b14e4f02588e1dc2 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:30:24 +0800 Subject: [PATCH 01/15] fix(kimi-code): flush wire journals before print-mode exit A print-mode turn's tail records (step.end / turn.ended / prompt.completed) are dispatched fire-and-forget and reach the journal only through the wire service's async persist queue. The print cleanup path never awaited that queue: with telemetry disabled (KIMI_DISABLE_TELEMETRY=1) cleanup returns in microseconds, and process.exit on the error path cut off the pending append-log flush, dropping the failed turn's closing records. Flush every session agent's dispatcher (which awaits the wire persist queue and the append-log store) before disposing the app, bounded by the shutdown timeout and best-effort so a persist failure never masks the turn outcome. --- .changeset/print-wire-flush-on-exit.md | 5 ++ apps/kimi-code/src/cli/v2/run-v2-print.ts | 30 +++++++ apps/kimi-code/test/cli/v2-run-print.test.ts | 87 ++++++++++++++++++++ 3 files changed, 122 insertions(+) create mode 100644 .changeset/print-wire-flush-on-exit.md diff --git a/.changeset/print-wire-flush-on-exit.md b/.changeset/print-wire-flush-on-exit.md new file mode 100644 index 0000000000..0ea49115d1 --- /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 a failed turn's closing session records on error exit. 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 f329c0e06c..c49ecba11f 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -30,6 +30,7 @@ import { IBootstrapService, IConfigService, IEventBus, + IEventDispatcher, IOAuthToolkit, ISessionIndex, ISessionManager, @@ -196,6 +197,7 @@ export async function runV2Print( } let restorePermission = async (): Promise => {}; + let flushWires = async (): Promise => {}; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; let telemetryService: ITelemetryService | undefined; @@ -205,6 +207,13 @@ export async function runV2Print( setCrashPhase('shutdown'); try { await restorePermission(); + // The turn's tail records (step.end / turn.ended / prompt.completed) are + // dispatched fire-and-forget and reach the journal only through the wire + // service's async persist queue. Without an explicit flush, a cleanup + // that returns fast (e.g. telemetry disabled) lets process.exit cut off + // that queue before the records land on disk. Best-effort: a persist + // failure was already reported where the append failed. + await raceWithTimeout(flushWires(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { if (telemetryService !== undefined) { await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS); @@ -254,6 +263,7 @@ export async function runV2Print( const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr); restorePermission = resolved.restorePermission; + flushWires = () => flushSessionWires(resolved.session, resolved.agent); telemetryService.setContext({ session_id: resolved.session.id, model: resolved.telemetryModel }); setTelemetryContext({ sessionId: resolved.session.id }); @@ -869,6 +879,26 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number { return count; } +/** + * Flush every session agent's wire journal. The main agent handle is included + * explicitly: `IAgentLifecycleService` skips `closing` agents, but a closing + * agent's already-dispatched tail records still deserve to land on disk. + */ +async function flushSessionWires( + session: ISessionScopeHandle, + mainAgent: IAgentScopeHandle, +): Promise { + 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); + } + await Promise.all( + [...handles].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 fdb114de2a..efddf2e135 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -17,6 +17,7 @@ import { IBootstrapService, IConfigService, IEventBus, + IEventDispatcher, IFileSystemStorageService, IOAuthToolkit, ISessionIndex, @@ -186,6 +187,7 @@ function makeFakeHarness() { [IAgentTaskService, { list: vi.fn(() => []) }], [IAgentCronService, { getNextFireTime: vi.fn(() => null) }], [IAgentGoalService, goal], + [IEventDispatcher, { flush: vi.fn(async () => {}) }], [ IAgentScopeContext, makeAgentScopeContext({ agentId: 'main', agentScope: 'agents/main' }), @@ -613,4 +615,89 @@ 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(); + }); }); From 463967f22671cd2a3510253d0e2fa794a1f75d5b Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:36:59 +0800 Subject: [PATCH 02/15] fix(kimi-code): settle each agent's wire flush independently Promise.all rejects as soon as one agent's flush fails; the best-effort caller then proceeds to app.dispose() and process.exit while the remaining agents' flushes are still in flight, re-exposing their journals to the truncation this cleanup is meant to prevent. Await every flush with Promise.allSettled instead (matches the drain helper's convention), and cover it with a two-agent regression test. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 10 ++-- apps/kimi-code/test/cli/v2-run-print.test.ts | 58 +++++++++++++++++++- 2 files changed, 63 insertions(+), 5 deletions(-) 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 c49ecba11f..377fa6ae3d 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -883,6 +883,8 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number { * Flush every session agent's wire journal. The main agent handle is included * explicitly: `IAgentLifecycleService` skips `closing` agents, but a closing * agent's already-dispatched tail records still deserve to land on disk. + * Each flush settles independently: one agent's broken journal must not cut + * the wait short for the others (the caller proceeds to `process.exit`). */ async function flushSessionWires( session: ISessionScopeHandle, @@ -894,7 +896,7 @@ async function flushSessionWires( const handle = agentManager.handleOf(agent.agentId); if (handle !== undefined) handles.add(handle); } - await Promise.all( + await Promise.allSettled( [...handles].map((handle) => handle.accessor.get(IEventDispatcher).flush()), ); } @@ -931,11 +933,11 @@ async function drainBackgroundTasks( allWaiters.push(waiter); } } - if (suppressions.length > 0) await Promise.all(suppressions); + if (suppressions.length > 0) await Promise.allSettled(suppressions); if (activeCount === 0 || batch.length === 0) break; - await Promise.all(batch); + await Promise.allSettled(batch); } - if (allWaiters.length > 0) await Promise.all(allWaiters); + if (allWaiters.length > 0) await Promise.allSettled(allWaiters); } function formatNativeTurnFailure(result: LoopRunResult): string { 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 efddf2e135..41a3f8455a 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -276,7 +276,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', () => { @@ -700,4 +700,60 @@ describe('runV2Print', () => { ); expect(app.dispose).toHaveBeenCalled(); }); + + it('keeps waiting for healthy agents when another agent\'s flush fails', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agent, agentServices, sessionServices } = 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 order: string[] = []; + const mainDispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + mainDispatcher.flush.mockRejectedValueOnce(new Error('disk full')); + const subAgent = fakeScope( + 'sub', + new Map([ + [ + IEventDispatcher, + { + flush: vi.fn(async () => { + await new Promise((resolve) => setTimeout(resolve, 10)); + order.push('sub-flushed'); + }), + }, + ], + ]), + ); + const lifecycle = sessionServices.get(IAgentLifecycleService) as { + list: ReturnType; + handleOf: ReturnType; + }; + lifecycle.list.mockReturnValue([{ agentId: 'main' }, { agentId: 'sub' }]); + lifecycle.handleOf.mockImplementation((id: string) => (id === 'sub' ? subAgent : agent)); + app.dispose.mockImplementation(() => { + order.push('disposed'); + }); + + 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(order).toEqual(['sub-flushed', 'disposed']); + }); }); From 84fe69b9247bc3932f120093ded895054fd489aa Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:45:32 +0800 Subject: [PATCH 03/15] fix(kimi-code): run print-mode shutdown phases concurrently MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wire flush, v2 telemetry shutdown, and v1 telemetry shutdown each hold a 3s allowance; run sequentially they can take 9s, past the 8s outer cleanup bound — the caller's process.exit would then cut off the tail (app.dispose included). The phases are independent, so await them concurrently (worst case one allowance) while keeping the v2 shutdown failure propagation and the best-effort semantics of the other two. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 29 ++++++++----- apps/kimi-code/test/cli/v2-run-print.test.ts | 45 ++++++++++++++++++++ 2 files changed, 63 insertions(+), 11 deletions(-) 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 377fa6ae3d..6bb2409984 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -207,18 +207,25 @@ export async function runV2Print( setCrashPhase('shutdown'); try { await restorePermission(); - // The turn's tail records (step.end / turn.ended / prompt.completed) are - // dispatched fire-and-forget and reach the journal only through the wire - // service's async persist queue. Without an explicit flush, a cleanup - // that returns fast (e.g. telemetry disabled) lets process.exit cut off - // that queue before the records land on disk. Best-effort: a persist - // failure was already reported where the append failed. - await raceWithTimeout(flushWires(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { - if (telemetryService !== undefined) { - await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS); - } - await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }).catch(() => {}); + // The shutdown phases are independent of each other; run them + // concurrently so their individual allowances cannot sum past the + // outer PROMPT_CLEANUP_TIMEOUT_MS bound and let the caller's + // process.exit cut off the tail (app.dispose included). + await Promise.all([ + // The turn's tail records (step.end / turn.ended / prompt.completed) + // are dispatched fire-and-forget and reach the journal only through + // the wire service's async persist queue. Without an explicit flush, + // a cleanup that returns fast (e.g. telemetry disabled) lets + // process.exit cut off that queue before the records land on disk. + // Best-effort: a persist failure was already reported where the + // append failed. + 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(); } })()); 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 41a3f8455a..345a97f3bd 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -756,4 +756,49 @@ describe('runV2Print', () => { ); expect(order).toEqual(['sub-flushed', 'disposed']); }); + + it('runs wire flush and telemetry shutdown concurrently', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agentServices, appServices } = makeFakeHarness(); + + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + let releaseFlush!: () => void; + const flushGate = new Promise((resolve) => { + releaseFlush = resolve; + }); + dispatcher.flush.mockReturnValueOnce(flushGate); + const telemetry = appServices.get(ITelemetryService) as { + shutdown: ReturnType; + }; + let releaseTelemetry!: () => void; + const telemetryGate = new Promise((resolve) => { + releaseTelemetry = resolve; + }); + telemetry.shutdown.mockReturnValueOnce(telemetryGate); + + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); + + const run = runV2Print(opts() as never, '1.2.3-test', { stdout, stderr }); + // A sequential cleanup would sit in the wire flush's 3s allowance before + // even starting telemetry shutdown; concurrent phases are both in flight + // well within that window. + for ( + let i = 0; + i < 200 && + (dispatcher.flush.mock.calls.length === 0 || telemetry.shutdown.mock.calls.length === 0); + i++ + ) { + await new Promise((resolve) => setTimeout(resolve, 10)); + } + expect(dispatcher.flush).toHaveBeenCalled(); + expect(telemetry.shutdown).toHaveBeenCalled(); + releaseFlush(); + releaseTelemetry(); + await run; + expect(app.dispose).toHaveBeenCalled(); + }); }); From f8fb8c8c87f3226cbedce19887521c7e53c3a66a Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:57:53 +0800 Subject: [PATCH 04/15] fix(kimi-code): quiesce active turns before the print-mode wire flush MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A termination signal (SIGINT/SIGTERM/SIGHUP) can arrive mid-turn: the cleanup ran the wire flush immediately, while the still-running turn only produced its cancellation and closing records from dispose()'s fire-and-forget teardown — after the flush, and after the signal handler's process.exit. Cancel every session agent's queued and active turns and await loop idleness before flushing (best-effort, bounded by the shutdown timeout); idle loops make this a no-op on the normal exit paths. Torn-down agent scopes are skipped defensively, mirroring the flush's per-agent settlement. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 63 +++++++++++++-- apps/kimi-code/test/cli/v2-run-print.test.ts | 80 ++++++++++++++++++++ 2 files changed, 136 insertions(+), 7 deletions(-) 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 6bb2409984..1291788720 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, @@ -197,6 +198,7 @@ export async function runV2Print( } let restorePermission = async (): Promise => {}; + let quiesceAgents = async (): Promise => {}; let flushWires = async (): Promise => {}; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; @@ -207,6 +209,11 @@ export async function runV2Print( setCrashPhase('shutdown'); try { await restorePermission(); + // A termination signal can arrive mid-turn: cancel any queued/active + // turns and let the loops go idle first, so the turn's cancellation + // and closing records exist before the flush below drains the + // journals. No-op when every loop is already idle. + await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { // The shutdown phases are independent of each other; run them // concurrently so their individual allowances cannot sum past the @@ -270,6 +277,7 @@ export async function runV2Print( const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr); restorePermission = resolved.restorePermission; + quiesceAgents = () => quiesceSessionAgents(resolved.session, resolved.agent); flushWires = () => flushSessionWires(resolved.session, resolved.agent); telemetryService.setContext({ session_id: resolved.session.id, model: resolved.telemetryModel }); @@ -887,24 +895,65 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number { } /** - * Flush every session agent's wire journal. The main agent handle is included + * Every agent handle in the session. The main agent handle is included * explicitly: `IAgentLifecycleService` skips `closing` agents, but a closing - * agent's already-dispatched tail records still deserve to land on disk. - * Each flush settles independently: one agent's broken journal must not cut - * the wait short for the others (the caller proceeds to `process.exit`). + * agent's in-flight turn and already-dispatched tail records still deserve to + * land on disk. */ -async function flushSessionWires( +function collectSessionAgentHandles( session: ISessionScopeHandle, mainAgent: IAgentScopeHandle, -): Promise { +): 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]; +} + +/** + * Cancel every session agent's queued and active turns, then wait for the + * loops to go idle. A termination signal can arrive mid-turn: without this, + * the turn's cancellation and closing records would only be produced by + * dispose()'s asynchronous teardown — after the wire flush, and after + * process.exit. Cancelling is a no-op on idle loops, so the normal + * (completed/failed turn) exit pays nothing here. + */ +async function quiesceSessionAgents( + session: ISessionScopeHandle, + mainAgent: IAgentScopeHandle, +): Promise { + const loops = collectSessionAgentHandles(session, mainAgent).flatMap((handle) => { + try { + return [handle.accessor.get(IAgentLoopService)]; + } catch { + // A torn-down agent scope has no loop to quiesce; the wire flush below + // still covers its already-dispatched records. + return []; + } + }); + 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())); +} + +/** + * Flush every session agent's wire journal. Each flush settles independently: + * one agent's broken journal must not cut the wait short for the others (the + * caller proceeds to `process.exit`). + */ +async function flushSessionWires( + session: ISessionScopeHandle, + mainAgent: IAgentScopeHandle, +): Promise { await Promise.allSettled( - [...handles].map((handle) => handle.accessor.get(IEventDispatcher).flush()), + collectSessionAgentHandles(session, mainAgent).map((handle) => + handle.accessor.get(IEventDispatcher).flush(), + ), ); } 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 345a97f3bd..459ac5a809 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, @@ -188,6 +189,14 @@ function makeFakeHarness() { [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 () => {}), + }, + ], [ IAgentScopeContext, makeAgentScopeContext({ agentId: 'main', agentScope: 'agents/main' }), @@ -801,4 +810,75 @@ describe('runV2Print', () => { await run; 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; + }; + loop.status.mockReturnValue({ state: 'running', pendingTurnIds: [] }); + loop.cancel.mockImplementation(() => { + order.push('cancel'); + return true; + }); + loop.settled = vi.fn(async () => { + order.push('settled'); + }); + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + dispatcher.flush = vi.fn(async () => { + order.push('flush'); + }); + + // A turn that is still in flight when the signal arrives. + const promptService = agentServices.get(IAgentPromptService) as { + enqueue: ReturnType; + }; + let settleTurn!: (result: unknown) => void; + promptService.enqueue.mockResolvedValueOnce({ + launched: Promise.resolve({ + id: 1, + result: new Promise((resolve) => { + settleTurn = resolve; + }), + }), + }); + + 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') }); + await onSigint(); + + expect(order).toEqual(['cancel', 'settled', 'flush', 'exit:130']); + expect(await outcome).toBeInstanceOf(Error); + }); }); From e566535306306becd682139aecfd9552415a4c1c Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:14:47 +0800 Subject: [PATCH 05/15] fix(kimi-code): drain prompts and await prompt completion in print quiesce Two follow-ups from review: - The loop settles (releaseActiveTurn) before the prompt-settle chain dispatches prompt.completed, so awaiting loop idleness alone let the wire flush race the final record. Quiesce now also drains each agent's prompt queue and awaits the tracked prompt completions (registered at enqueue time) before the flush runs. - Restore drainBackgroundTasks to Promise.all: a previous edit unintentionally switched it to allSettled alongside the per-agent wire flush, silently discarding persistence failures from suppressTerminalNotification()/wait(). Only the independent wire flushes are meant to be best-effort. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 67 +++++++++++++++----- apps/kimi-code/test/cli/v2-run-print.test.ts | 21 +++++- 2 files changed, 70 insertions(+), 18 deletions(-) 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 1291788720..ffe3b7cf19 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -200,6 +200,7 @@ export async function runV2Print( let restorePermission = async (): Promise => {}; let quiesceAgents = async (): Promise => {}; let flushWires = async (): Promise => {}; + const promptCompletions: Promise[] = []; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; let telemetryService: ITelemetryService | undefined; @@ -209,10 +210,11 @@ export async function runV2Print( setCrashPhase('shutdown'); try { await restorePermission(); - // A termination signal can arrive mid-turn: cancel any queued/active - // turns and let the loops go idle first, so the turn's cancellation - // and closing records exist before the flush below drains the - // journals. No-op when every loop is already idle. + // A termination signal can arrive mid-turn: drain pending prompts, + // cancel any queued/active turns, and let the loops and prompt + // completions settle first, so the turn's cancellation and closing + // records exist before the flush below drains the journals. No-op + // when every loop is already idle. await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { // The shutdown phases are independent of each other; run them @@ -277,7 +279,10 @@ export async function runV2Print( const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr); restorePermission = resolved.restorePermission; - quiesceAgents = () => quiesceSessionAgents(resolved.session, resolved.agent); + quiesceAgents = async () => { + await quiesceSessionAgents(resolved.session, resolved.agent); + await Promise.allSettled(promptCompletions); + }; flushWires = () => flushSessionWires(resolved.session, resolved.agent); telemetryService.setContext({ session_id: resolved.session.id, model: resolved.telemetryModel }); @@ -288,6 +293,9 @@ export async function runV2Print( telemetryService.track2('first_launch'); } + const trackPromptCompletion = (completion: Promise): void => { + promptCompletions.push(completion); + }; const goalCreate = parseHeadlessGoalCreate(opts.prompt!); if (goalCreate !== undefined) { await runNativeGoal( @@ -299,6 +307,7 @@ export async function runV2Print( outputFormat, stdout, stderr, + trackPromptCompletion, ); } else { await runNativeTurn( @@ -309,6 +318,7 @@ export async function runV2Print( outputFormat, stdout, stderr, + trackPromptCompletion, ); } writeResumeHint(resolved.session.id, outputFormat, stdout, stderr); @@ -483,6 +493,7 @@ async function runNativeTurn( outputFormat: PromptOutputFormat, stdout: PromptOutput, stderr: PromptOutput, + trackPromptCompletion?: (completion: Promise) => void, ): Promise { const writer: PromptTurnWriter = outputFormat === 'stream-json' @@ -508,6 +519,11 @@ async function runNativeTurn( origin: { kind: 'user' }, }, }); + // The cleanup's quiesce phase awaits every registered completion before + // flushing: `prompt.completed` is dispatched from the prompt-settle chain + // that runs after the loop reports idle, so loop idleness alone does not + // guarantee the record exists yet. + trackPromptCompletion?.(handle.completion); const turn = await handle.launched; if (turn === undefined) { // A prompt blocked by an onBeforeSubmitPrompt hook never launches a turn. @@ -581,6 +597,7 @@ async function runNativeGoal( outputFormat: PromptOutputFormat, stdout: PromptOutput, stderr: PromptOutput, + trackPromptCompletion?: (completion: Promise) => void, ): Promise { requireConfiguredModel(model); const goalService = agent.accessor.get(IAgentGoalService); @@ -598,7 +615,16 @@ async function runNativeGoal( } }); try { - await runNativeTurn(app, session, agent, goal.objective, outputFormat, stdout, stderr); + await runNativeTurn( + app, + session, + agent, + goal.objective, + outputFormat, + stdout, + stderr, + trackPromptCompletion, + ); } finally { subscription.dispose(); const snapshot = completedSnapshot ?? goalService.getGoal().goal; @@ -914,18 +940,27 @@ function collectSessionAgentHandles( } /** - * Cancel every session agent's queued and active turns, then wait for the - * loops to go idle. A termination signal can arrive mid-turn: without this, - * the turn's cancellation and closing records would only be produced by - * dispose()'s asynchronous teardown — after the wire flush, and after - * process.exit. Cancelling is a no-op on idle loops, so the normal - * (completed/failed turn) exit pays nothing here. + * Drain every session agent's prompt queue, cancel queued and active turns, + * then wait for the loops to go idle. A termination signal can arrive + * mid-turn: without this, the turn's cancellation and closing records would + * only be produced by dispose()'s asynchronous teardown — after the wire + * flush, and after process.exit. Draining and cancelling are no-ops on idle + * agents, so the normal (completed/failed turn) exit pays nothing here. */ async function quiesceSessionAgents( session: ISessionScopeHandle, mainAgent: IAgentScopeHandle, ): Promise { - const loops = collectSessionAgentHandles(session, mainAgent).flatMap((handle) => { + const handles = collectSessionAgentHandles(session, mainAgent); + for (const handle of handles) { + try { + await handle.accessor.get(IAgentPromptService).drain(); + } catch { + // A torn-down agent scope has no prompt service to drain; the loop + // cancellation below still applies. + } + } + const loops = handles.flatMap((handle) => { try { return [handle.accessor.get(IAgentLoopService)]; } catch { @@ -989,11 +1024,11 @@ async function drainBackgroundTasks( allWaiters.push(waiter); } } - if (suppressions.length > 0) await Promise.allSettled(suppressions); + if (suppressions.length > 0) await Promise.all(suppressions); if (activeCount === 0 || batch.length === 0) break; - await Promise.allSettled(batch); + await Promise.all(batch); } - if (allWaiters.length > 0) await Promise.allSettled(allWaiters); + if (allWaiters.length > 0) await Promise.all(allWaiters); } function formatNativeTurnFailure(result: LoopRunResult): string { 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 459ac5a809..ef37c35ae7 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -183,6 +183,7 @@ function makeFakeHarness() { }), }; }), + drain: vi.fn(async () => {}), }, ], [IAgentTaskService, { list: vi.fn(() => []) }], @@ -837,11 +838,14 @@ describe('runV2Print', () => { order.push('flush'); }); - // A turn that is still in flight when the signal arrives. + // A turn that is still in flight when the signal arrives; its prompt + // completion stays pending until released below. const promptService = agentServices.get(IAgentPromptService) as { enqueue: ReturnType; + drain: ReturnType; }; let settleTurn!: (result: unknown) => void; + let releaseCompletion!: () => void; promptService.enqueue.mockResolvedValueOnce({ launched: Promise.resolve({ id: 1, @@ -849,6 +853,9 @@ describe('runV2Print', () => { settleTurn = resolve; }), }), + completion: new Promise((resolve) => { + releaseCompletion = resolve; + }), }); const handlers = new Map Promise>(); @@ -876,7 +883,17 @@ describe('runV2Print', () => { } const onSigint = handlers.get('SIGINT')!; settleTurn({ type: 'cancelled', steps: 0, reason: new Error('aborted') }); - await onSigint(); + const sigintRun = onSigint(); + // Quiesce awaits the prompt completion before the wire flush: while the + // completion is still pending, the loops are already idle but the flush + // must not have happened. + for (let i = 0; i < 100 && !order.includes('settled'); i++) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(promptService.drain).toHaveBeenCalled(); + expect(order).toEqual(['cancel', 'settled']); + releaseCompletion(); + await sigintRun; expect(order).toEqual(['cancel', 'settled', 'flush', 'exit:130']); expect(await outcome).toBeInstanceOf(Error); From 61e34b089fae770fd186fdb5408046d4ca7badeb Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:28:09 +0800 Subject: [PATCH 06/15] fix(kimi-code): cover every agent's prompt queue in print quiesce MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tracking only the main run's prompt completion left subagent and background prompts racing the wire flush the same way: the loop reports idle before that agent's prompt-settle chain dispatches prompt.completed. Replace the per-handle completion tracking with a uniform wait — after draining and cancelling, poll every agent's prompt queue snapshot until none reports an active or pending prompt. settle() clears the active prompt and dispatches the record in one synchronous block, so an empty snapshot proves the record is already queued for the flush. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 79 ++++++++++---------- apps/kimi-code/test/cli/v2-run-print.test.ts | 28 ++++--- 2 files changed, 56 insertions(+), 51 deletions(-) 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 ffe3b7cf19..a703326a12 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -124,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 @@ -200,7 +202,6 @@ export async function runV2Print( let restorePermission = async (): Promise => {}; let quiesceAgents = async (): Promise => {}; let flushWires = async (): Promise => {}; - const promptCompletions: Promise[] = []; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; let telemetryService: ITelemetryService | undefined; @@ -211,10 +212,10 @@ export async function runV2Print( try { await restorePermission(); // A termination signal can arrive mid-turn: drain pending prompts, - // cancel any queued/active turns, and let the loops and prompt - // completions settle first, so the turn's cancellation and closing - // records exist before the flush below drains the journals. No-op - // when every loop is already idle. + // cancel any queued/active turns, and wait for every agent's loops + // and prompt queues to quiesce first, so the turn's cancellation and + // closing records exist before the flush below drains the journals. + // No-op when every agent is already idle. await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { // The shutdown phases are independent of each other; run them @@ -279,10 +280,7 @@ export async function runV2Print( const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr); restorePermission = resolved.restorePermission; - quiesceAgents = async () => { - await quiesceSessionAgents(resolved.session, resolved.agent); - await Promise.allSettled(promptCompletions); - }; + quiesceAgents = () => quiesceSessionAgents(resolved.session, resolved.agent); flushWires = () => flushSessionWires(resolved.session, resolved.agent); telemetryService.setContext({ session_id: resolved.session.id, model: resolved.telemetryModel }); @@ -293,9 +291,6 @@ export async function runV2Print( telemetryService.track2('first_launch'); } - const trackPromptCompletion = (completion: Promise): void => { - promptCompletions.push(completion); - }; const goalCreate = parseHeadlessGoalCreate(opts.prompt!); if (goalCreate !== undefined) { await runNativeGoal( @@ -307,7 +302,6 @@ export async function runV2Print( outputFormat, stdout, stderr, - trackPromptCompletion, ); } else { await runNativeTurn( @@ -318,7 +312,6 @@ export async function runV2Print( outputFormat, stdout, stderr, - trackPromptCompletion, ); } writeResumeHint(resolved.session.id, outputFormat, stdout, stderr); @@ -493,7 +486,6 @@ async function runNativeTurn( outputFormat: PromptOutputFormat, stdout: PromptOutput, stderr: PromptOutput, - trackPromptCompletion?: (completion: Promise) => void, ): Promise { const writer: PromptTurnWriter = outputFormat === 'stream-json' @@ -519,11 +511,6 @@ async function runNativeTurn( origin: { kind: 'user' }, }, }); - // The cleanup's quiesce phase awaits every registered completion before - // flushing: `prompt.completed` is dispatched from the prompt-settle chain - // that runs after the loop reports idle, so loop idleness alone does not - // guarantee the record exists yet. - trackPromptCompletion?.(handle.completion); const turn = await handle.launched; if (turn === undefined) { // A prompt blocked by an onBeforeSubmitPrompt hook never launches a turn. @@ -597,7 +584,6 @@ async function runNativeGoal( outputFormat: PromptOutputFormat, stdout: PromptOutput, stderr: PromptOutput, - trackPromptCompletion?: (completion: Promise) => void, ): Promise { requireConfiguredModel(model); const goalService = agent.accessor.get(IAgentGoalService); @@ -615,16 +601,7 @@ async function runNativeGoal( } }); try { - await runNativeTurn( - app, - session, - agent, - goal.objective, - outputFormat, - stdout, - stderr, - trackPromptCompletion, - ); + await runNativeTurn(app, session, agent, goal.objective, outputFormat, stdout, stderr); } finally { subscription.dispose(); const snapshot = completedSnapshot ?? goalService.getGoal().goal; @@ -941,25 +918,27 @@ function collectSessionAgentHandles( /** * Drain every session agent's prompt queue, cancel queued and active turns, - * then wait for the loops to go idle. A termination signal can arrive - * mid-turn: without this, the turn's cancellation and closing records would - * only be produced by dispose()'s asynchronous teardown — after the wire - * flush, and after process.exit. Draining and cancelling are no-ops on idle - * agents, so the normal (completed/failed turn) exit pays nothing here. + * then wait for the loops to go idle and the prompt queues to empty. A + * termination signal can arrive mid-turn: without this, the turn's + * cancellation and closing records would only be produced by dispose()'s + * asynchronous teardown — after the wire flush, and after process.exit. + * Draining and cancelling are no-ops on idle agents, so the normal + * (completed/failed turn) exit pays nothing here. */ async function quiesceSessionAgents( session: ISessionScopeHandle, mainAgent: IAgentScopeHandle, ): Promise { const handles = collectSessionAgentHandles(session, mainAgent); - for (const handle of handles) { + const promptServices = handles.flatMap((handle) => { try { - await handle.accessor.get(IAgentPromptService).drain(); + return [handle.accessor.get(IAgentPromptService)]; } catch { - // A torn-down agent scope has no prompt service to drain; the loop - // cancellation below still applies. + // A torn-down agent scope has no prompt service to drain or observe. + return []; } - } + }); + await Promise.allSettled(promptServices.map((service) => service.drain())); const loops = handles.flatMap((handle) => { try { return [handle.accessor.get(IAgentLoopService)]; @@ -974,6 +953,24 @@ async function quiesceSessionAgents( loop.cancel(); } await Promise.allSettled(loops.map((loop) => loop.settled())); + // The loop reports idle (releaseActiveTurn) before the prompt-settle chain + // dispatches the final prompt.completed; settle() clears the active prompt + // and dispatches the record in one synchronous block, so an empty queue + // snapshot proves the record was already queued for the wire flush. + for (;;) { + const busy = promptServices.some((service) => { + try { + const snapshot = service.list(); + return snapshot.active !== undefined || snapshot.pending.length > 0; + } catch { + return false; + } + }); + if (!busy) return; + await new Promise((resolve) => { + setTimeout(resolve, PROMPT_QUIESCE_POLL_MS); + }); + } } /** 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 ef37c35ae7..38d62bddb6 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -184,6 +184,7 @@ function makeFakeHarness() { }; }), drain: vi.fn(async () => {}), + list: vi.fn(() => ({ active: undefined, pending: [] })), }, ], [IAgentTaskService, { list: vi.fn(() => []) }], @@ -838,14 +839,15 @@ describe('runV2Print', () => { order.push('flush'); }); - // A turn that is still in flight when the signal arrives; its prompt - // completion stays pending until released below. + // A turn that is still in flight when the signal arrives; its prompt queue + // stays non-empty (the prompt-settle chain has not run) until released + // below. const promptService = agentServices.get(IAgentPromptService) as { enqueue: ReturnType; drain: ReturnType; + list: ReturnType; }; let settleTurn!: (result: unknown) => void; - let releaseCompletion!: () => void; promptService.enqueue.mockResolvedValueOnce({ launched: Promise.resolve({ id: 1, @@ -853,10 +855,13 @@ describe('runV2Print', () => { settleTurn = resolve; }), }), - completion: new Promise((resolve) => { - releaseCompletion = resolve; - }), }); + let promptsBusy = true; + promptService.list = vi.fn(() => + promptsBusy + ? { active: { id: 'p1' }, pending: [] } + : { active: undefined, pending: [] }, + ); const handlers = new Map Promise>(); const fakeProcess = { @@ -884,15 +889,18 @@ describe('runV2Print', () => { const onSigint = handlers.get('SIGINT')!; settleTurn({ type: 'cancelled', steps: 0, reason: new Error('aborted') }); const sigintRun = onSigint(); - // Quiesce awaits the prompt completion before the wire flush: while the - // completion is still pending, the loops are already idle but the flush - // must not have happened. + // Quiesce waits for the prompt queue to empty before the wire flush: the + // loops are already idle, but while the snapshot still reports an active + // prompt the flush must not happen. for (let i = 0; i < 100 && !order.includes('settled'); i++) { await new Promise((resolve) => setTimeout(resolve, 5)); } + // A few quiesce poll cycles (the poll interval is 10ms): enough time for + // the loop to have flushed if it were not actually waiting on the queue. + await new Promise((resolve) => setTimeout(resolve, 30)); expect(promptService.drain).toHaveBeenCalled(); expect(order).toEqual(['cancel', 'settled']); - releaseCompletion(); + promptsBusy = false; await sigintRun; expect(order).toEqual(['cancel', 'settled', 'flush', 'exit:130']); From a755e738b89a2acdec85ca6fba99c2f774b15ae7 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:48:41 +0800 Subject: [PATCH 07/15] fix(agent-core-v2,kimi-code): cover the prompt launch window in print quiesce A termination signal can also arrive while startNext() is mid-launch: the prompt has left pending and is not active yet, so drain() cannot cancel it and an active/pending snapshot reads empty. Quiesce would then flush and exit while the launch still dispatches records afterward. Expose the service's launching phase on PromptQueueSnapshot and make the print cleanup's quiesce a repeat-until-idle loop: every pass drains, cancels, and awaits the loops, then re-checks launching/active/pending, so a prompt surfacing from the launch window is cancelled on the next pass instead of escaping the flush. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 30 +++++++++------ apps/kimi-code/test/cli/v2-run-print.test.ts | 38 +++++++++++-------- .../agent-core-v2/src/agent/prompt/prompt.ts | 1 + .../src/agent/prompt/promptService.ts | 2 +- .../test/agent/prompt/promptService.test.ts | 24 +++++++++++- .../test/agent/undo/undo.test.ts | 1 + .../test/app/gateway/gateway.test.ts | 2 +- .../agentTitlePromptSourceService.test.ts | 5 ++- 8 files changed, 72 insertions(+), 31 deletions(-) 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 a703326a12..caf0af47e8 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -938,7 +938,6 @@ async function quiesceSessionAgents( return []; } }); - await Promise.allSettled(promptServices.map((service) => service.drain())); const loops = handles.flatMap((handle) => { try { return [handle.accessor.get(IAgentLoopService)]; @@ -948,20 +947,29 @@ async function quiesceSessionAgents( return []; } }); - 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())); - // The loop reports idle (releaseActiveTurn) before the prompt-settle chain - // dispatches the final prompt.completed; settle() clears the active prompt - // and dispatches the record in one synchronous block, so an empty queue - // snapshot proves the record was already queued for the wire flush. + // Repeat until a full pass finds every queue empty: a prompt can surface + // after one pass already ran — from startNext's launch window (the prompt + // left `pending` and is not yet `active`, so drain could not see it) or + // from a cancelled turn's settle chain. The snapshot covers all three + // states; the loop reports idle (releaseActiveTurn) before the + // prompt-settle chain dispatches the final record, and settle() clears the + // active prompt and dispatches that record in one synchronous block, so an + // empty snapshot proves it was already queued for the wire flush. 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 busy = promptServices.some((service) => { try { const snapshot = service.list(); - return snapshot.active !== undefined || snapshot.pending.length > 0; + return ( + snapshot.launching || + snapshot.active !== undefined || + snapshot.pending.length > 0 + ); } catch { return false; } 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 38d62bddb6..106a55c8aa 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -184,7 +184,7 @@ function makeFakeHarness() { }; }), drain: vi.fn(async () => {}), - list: vi.fn(() => ({ active: undefined, pending: [] })), + list: vi.fn(() => ({ launching: false, active: undefined, pending: [] })), }, ], [IAgentTaskService, { list: vi.fn(() => []) }], @@ -826,11 +826,11 @@ describe('runV2Print', () => { }; loop.status.mockReturnValue({ state: 'running', pendingTurnIds: [] }); loop.cancel.mockImplementation(() => { - order.push('cancel'); + if (!order.includes('cancel')) order.push('cancel'); return true; }); loop.settled = vi.fn(async () => { - order.push('settled'); + if (!order.includes('settled')) order.push('settled'); }); const dispatcher = agentServices.get(IEventDispatcher) as { flush: ReturnType; @@ -839,9 +839,10 @@ describe('runV2Print', () => { order.push('flush'); }); - // A turn that is still in flight when the signal arrives; its prompt queue - // stays non-empty (the prompt-settle chain has not run) until released - // below. + // A turn that is still in flight when the signal arrives. The prompt + // queue reports the launch window first (the prompt left `pending` and is + // not yet `active`), then the running prompt, and only goes empty once + // released below. const promptService = agentServices.get(IAgentPromptService) as { enqueue: ReturnType; drain: ReturnType; @@ -856,12 +857,16 @@ describe('runV2Print', () => { }), }), }); - let promptsBusy = true; - promptService.list = vi.fn(() => - promptsBusy - ? { active: { id: 'p1' }, pending: [] } - : { active: undefined, pending: [] }, - ); + 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 = { @@ -890,8 +895,8 @@ describe('runV2Print', () => { settleTurn({ type: 'cancelled', steps: 0, reason: new Error('aborted') }); const sigintRun = onSigint(); // Quiesce waits for the prompt queue to empty before the wire flush: the - // loops are already idle, but while the snapshot still reports an active - // prompt the flush must not happen. + // loops are already idle, but while the snapshot still reports the launch + // window or an active prompt the flush must not happen. for (let i = 0; i < 100 && !order.includes('settled'); i++) { await new Promise((resolve) => setTimeout(resolve, 5)); } @@ -900,7 +905,10 @@ describe('runV2Print', () => { await new Promise((resolve) => setTimeout(resolve, 30)); expect(promptService.drain).toHaveBeenCalled(); expect(order).toEqual(['cancel', 'settled']); - promptsBusy = false; + promptPhase = 'active'; + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(order).toEqual(['cancel', 'settled']); + promptPhase = 'empty'; await sigintRun; expect(order).toEqual(['cancel', 'settled', 'flush', 'exit:130']); diff --git a/packages/agent-core-v2/src/agent/prompt/prompt.ts b/packages/agent-core-v2/src/agent/prompt/prompt.ts index 4eb7136d2a..fc17432448 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 81dcac6947..90a2c42193 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/test/agent/prompt/promptService.test.ts b/packages/agent-core-v2/test/agent/prompt/promptService.test.ts index c1eece1b3f..3733283b4b 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/undo/undo.test.ts b/packages/agent-core-v2/test/agent/undo/undo.test.ts index d88b161d42..925b16d39c 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 c015d127d8..01d1ecbc50 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/sessionTitle/agentTitlePromptSourceService.test.ts b/packages/agent-core-v2/test/session/sessionTitle/agentTitlePromptSourceService.test.ts index 0c5bd74722..182ebb3e73 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', From 0614a7abb9c9bb9cd5f6f23e0cf6fb5717697935 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:01:37 +0800 Subject: [PATCH 08/15] fix(kimi-code): freeze loop producers across the print-mode flush MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An empty prompt queue is only a point-in-time observation: background task completions and cron fires enqueue straight into the loop, bypassing the prompt queue, so a late producer could still start a turn — and new records — after the journals were flushed. Quiesce now leaves a quiescence guard held on every loop once queues read empty; the caller holds the release across the wire flush and app disposal, so late submissions queue behind the guard and are rejected by disposal instead of racing process.exit. Loops that refuse the guard (still busy) send the quiesce into another drain/cancel/settle pass. --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 93 ++++++++++++++------ apps/kimi-code/test/cli/v2-run-print.test.ts | 39 ++++++++ 2 files changed, 105 insertions(+), 27 deletions(-) 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 caf0af47e8..2132cbc248 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -201,6 +201,7 @@ 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; @@ -218,25 +219,32 @@ export async function runV2Print( // No-op when every agent is already idle. await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { - // The shutdown phases are independent of each other; run them - // concurrently so their individual allowances cannot sum past the - // outer PROMPT_CLEANUP_TIMEOUT_MS bound and let the caller's - // process.exit cut off the tail (app.dispose included). - await Promise.all([ - // The turn's tail records (step.end / turn.ended / prompt.completed) - // are dispatched fire-and-forget and reach the journal only through - // the wire service's async persist queue. Without an explicit flush, - // a cleanup that returns fast (e.g. telemetry disabled) lets - // process.exit cut off that queue before the records land on disk. - // Best-effort: a persist failure was already reported where the - // append failed. - 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(); + try { + // The shutdown phases are independent of each other; run them + // concurrently so their individual allowances cannot sum past the + // outer PROMPT_CLEANUP_TIMEOUT_MS bound and let the caller's + // process.exit cut off the tail (app.dispose included). + await Promise.all([ + // The turn's tail records (step.end / turn.ended / prompt.completed) + // are dispatched fire-and-forget and reach the journal only through + // the wire service's async persist queue. Without an explicit flush, + // a cleanup that returns fast (e.g. telemetry disabled) lets + // process.exit cut off that queue before the records land on disk. + // Best-effort: a persist failure was already reported where the + // append failed. + 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 { + // Quiescence guards stay held across the flush and disposal so a + // late background-task completion or cron fire cannot start a new + // turn (and new records) after the journals were drained. + releaseQuiescence?.(); + } } })()); await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS); @@ -280,7 +288,9 @@ export async function runV2Print( const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr); restorePermission = resolved.restorePermission; - quiesceAgents = () => quiesceSessionAgents(resolved.session, resolved.agent); + 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 }); @@ -924,11 +934,19 @@ function collectSessionAgentHandles( * asynchronous teardown — after the wire flush, and after process.exit. * Draining and cancelling are no-ops on idle agents, so the normal * (completed/failed turn) exit pays nothing here. + * + * On success returns a release function: every loop is left holding a + * quiescence guard so late producers (a background-task completion or a cron + * fire, both of which enqueue straight into the loop, bypassing the prompt + * queue) cannot start a new turn — and new records — after the prompt + * queues read empty. The caller holds the release across the wire flush and + * disposal. Returns undefined when quiescence could not be reached within + * the caller's bound. */ async function quiesceSessionAgents( session: ISessionScopeHandle, mainAgent: IAgentScopeHandle, -): Promise { +): Promise<(() => void) | undefined> { const handles = collectSessionAgentHandles(session, mainAgent); const promptServices = handles.flatMap((handle) => { try { @@ -947,11 +965,11 @@ async function quiesceSessionAgents( return []; } }); - // Repeat until a full pass finds every queue empty: a prompt can surface - // after one pass already ran — from startNext's launch window (the prompt - // left `pending` and is not yet `active`, so drain could not see it) or - // from a cancelled turn's settle chain. The snapshot covers all three - // states; the loop reports idle (releaseActiveTurn) before the + // Repeat until a full pass finds every queue empty and every loop freezable: + // a prompt can surface after one pass already ran — from startNext's launch + // window (the prompt left `pending` and is not yet `active`, so drain could + // not see it) or from a cancelled turn's settle chain. The snapshot covers + // all three states; the loop reports idle (releaseActiveTurn) before the // prompt-settle chain dispatches the final record, and settle() clears the // active prompt and dispatches that record in one synchronous block, so an // empty snapshot proves it was already queued for the wire flush. @@ -962,6 +980,22 @@ async function quiesceSessionAgents( 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(); @@ -974,7 +1008,12 @@ async function quiesceSessionAgents( return false; } }); - if (!busy) return; + 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); }); 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 106a55c8aa..c139615a65 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -197,6 +197,7 @@ function makeFakeHarness() { status: vi.fn(() => ({ state: 'idle', pendingTurnIds: [] })), cancel: vi.fn(() => false), settled: vi.fn(async () => {}), + tryAcquireQuiescence: vi.fn(() => ({ dispose: vi.fn() })), }, ], [ @@ -823,6 +824,7 @@ describe('runV2Print', () => { status: ReturnType; cancel: ReturnType; settled: ReturnType; + tryAcquireQuiescence: ReturnType; }; loop.status.mockReturnValue({ state: 'running', pendingTurnIds: [] }); loop.cancel.mockImplementation(() => { @@ -832,6 +834,8 @@ describe('runV2Print', () => { loop.settled = vi.fn(async () => { if (!order.includes('settled')) order.push('settled'); }); + const guardDispose = vi.fn(); + loop.tryAcquireQuiescence = vi.fn(() => ({ dispose: guardDispose })); const dispatcher = agentServices.get(IEventDispatcher) as { flush: ReturnType; }; @@ -912,6 +916,41 @@ describe('runV2Print', () => { await sigintRun; expect(order).toEqual(['cancel', 'settled', 'flush', 'exit:130']); + // Producers stay frozen across the flush and disposal: 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); }); + + it('retries quiescence acquisition while a producer keeps the loop busy', async () => { + const stdout = writer(); + const stderr = writer(); + const { app, agentServices } = makeFakeHarness(); + + const loop = agentServices.get(IAgentLoopService) as { + tryAcquireQuiescence: ReturnType; + }; + let attempts = 0; + loop.tryAcquireQuiescence = vi.fn(() => { + attempts += 1; + // The first pass finds the loop busy (a background-task completion just + // enqueued a turn); the second pass can freeze it. + return attempts === 1 ? undefined : { dispose: vi.fn() }; + }); + + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); + + await runV2Print(opts() as never, '1.2.3-test', { stdout, stderr }); + + expect(attempts).toBeGreaterThanOrEqual(2); + const dispatcher = agentServices.get(IEventDispatcher) as { + flush: ReturnType; + }; + expect(dispatcher.flush).toHaveBeenCalled(); + expect(app.dispose).toHaveBeenCalled(); + }); }); From dda1e054eb711291d2fb73af87ebf5bda0276664 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:10:43 +0800 Subject: [PATCH 09/15] fix(agent-core-v2): flush the wire journal when removing an agent AgentLifecycleService.remove() quiesced the agent but never flushed its wire journal: the records reached the append log only through the fire-and-forget retirement flush, so a process exiting right after a subagent's removal (e.g. print mode on a termination signal) could truncate the closing records. Flush the agent's event dispatcher after the quiesce and before disposal; a persist failure is reported without blocking the removal. --- .../src/session/agentLifecycle/agentLifecycleService.ts | 2 ++ .../test/session/agentLifecycle/agentLifecycle.test.ts | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index 6afd674621..aa42be3a43 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 { @@ -373,6 +374,7 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle compaction.abortController.abort(reason); } await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]); + await handle.accessor.get(IEventDispatcher).flush().catch(onUnexpectedError); managed.killSpace(); await handle.dispose(); if (this.roster.get(agent.agentId) === managed) this.roster.delete(agent.agentId); 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 f8895d79ce..aadab02761 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -509,6 +509,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); From 5c739e50c01a00de0e8e70f56da2daf5ee336c50 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:25:33 +0800 Subject: [PATCH 10/15] fix(agent-core-v2): wait for the prompt queue to go idle before flushing a removed agent The remove() flush could still snapshot the persistence queue before the cancelled prompt's final record was appended: loop.settled() resolves in releaseActiveTurn() before the prompt-settle chain dispatches prompt.aborted/prompt.completed, and prompt.drain() does not await that settlement. Wait (bounded) for the prompt queue to report no launching, active, or pending prompts before flushing; an unreadable snapshot counts as idle so a wedged service cannot stall the removal. --- .../agentLifecycle/agentLifecycleService.ts | 16 +++++++++++++ .../agentLifecycle/agentLifecycle.test.ts | 23 +++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index aa42be3a43..30d67011d4 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -59,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(); @@ -374,6 +377,19 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle compaction.abortController.abort(reason); } await Promise.all([loop.settled(), compactionSettled, prompt.drain(reason)]); + const promptIdleDeadline = Date.now() + REMOVE_PROMPT_QUIESCE_TIMEOUT_MS; + for (;;) { + let idle = true; + try { + const snapshot = prompt.list(); + idle = + !snapshot.launching && snapshot.active === undefined && snapshot.pending.length === 0; + } catch { + idle = true; + } + if (idle || Date.now() >= promptIdleDeadline) break; + await new Promise((resolve) => setTimeout(resolve, REMOVE_PROMPT_QUIESCE_POLL_MS)); + } await handle.accessor.get(IEventDispatcher).flush().catch(onUnexpectedError); managed.killSpace(); await handle.dispose(); 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 aadab02761..1843aeb4bd 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -353,6 +353,7 @@ describe('AgentLifecycleService', () => { ix.stub(IAgentPromptService, { _serviceBrand: undefined, drain: promptDrain, + list: () => ({ launching: false, active: undefined, pending: [] }), } as unknown as IAgentPromptService); ix.stub(ITelemetryService, { _serviceBrand: undefined, @@ -518,6 +519,28 @@ describe('AgentLifecycleService', () => { expect(flush).toHaveBeenCalled(); }); + it('remove waits for the prompt queue to go idle before flushing', async () => { + const svc = ix.get(IAgentLifecycleService); + await svc.create({ agentId: 'main' }); + const promptService = svc.handleOf('main')!.accessor.get(IAgentPromptService) as unknown as { + list: ReturnType; + }; + let queueIdle = false; + promptService.list = vi.fn(() => + queueIdle + ? { launching: false, active: undefined, pending: [] } + : { launching: false, active: { id: 'p1' }, pending: [] }, + ); + const dispatcher = svc.handleOf('main')!.accessor.get(IEventDispatcher); + const flush = vi.spyOn(dispatcher, 'flush'); + const removal = svc.remove(svc.get('main')!); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(flush).not.toHaveBeenCalled(); + queueIdle = true; + await removal; + 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); From e3aa0a6a42616f69069924a914f6bb1a3418e7ff Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:26:58 +0800 Subject: [PATCH 11/15] chore: widen the print wire-flush changeset to termination signals --- .changeset/print-wire-flush-on-exit.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.changeset/print-wire-flush-on-exit.md b/.changeset/print-wire-flush-on-exit.md index 0ea49115d1..6e649d403b 100644 --- a/.changeset/print-wire-flush-on-exit.md +++ b/.changeset/print-wire-flush-on-exit.md @@ -2,4 +2,4 @@ "@moonshot-ai/kimi-code": patch --- -Fix print mode (`kimi -p`) losing a failed turn's closing session records on error exit. +Fix print mode (`kimi -p`) losing session records when the run exits on an error or a termination signal. From f4a76a443367d85d8c5b2ab2441775f295bc6bd7 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:32:59 +0800 Subject: [PATCH 12/15] fix(agent-core-v2): re-cancel prompts that finish launching during removal A prompt mid-launch (startNext awaiting daemon materialization or a pre-submit hook) is invisible to drain() and has no turn to cancel yet; the passive idle wait would then let it run to the deadline and flush before its closing records existed. Repeat drain/cancel/settle on every pass until the queue reads idle, and once idle hold a quiescence guard across the flush and disposal so a late producer cannot start new work in between. --- .../agentLifecycle/agentLifecycleService.ts | 34 ++++++++++---- .../agentLifecycle/agentLifecycle.test.ts | 47 +++++++++++++++++++ 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index 30d67011d4..635a4ecdff 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -369,16 +369,17 @@ 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)]); 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(); @@ -387,12 +388,27 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle } catch { idle = true; } - if (idle || Date.now() >= promptIdleDeadline) break; + 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)); } - await handle.accessor.get(IEventDispatcher).flush().catch(onUnexpectedError); - managed.killSpace(); - await handle.dispose(); + 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/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index 1843aeb4bd..6567041425 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -348,6 +348,7 @@ describe('AgentLifecycleService', () => { }), cancel: loopCancel, settled: loopSettled, + tryAcquireQuiescence: vi.fn(() => ({ dispose: vi.fn() })), } as unknown as IAgentLoopService); promptDrain = vi.fn(async () => {}); ix.stub(IAgentPromptService, { @@ -541,6 +542,52 @@ describe('AgentLifecycleService', () => { expect(flush).toHaveBeenCalled(); }); + it('remove re-cancels a prompt that finishes launching during the removal', async () => { + const svc = ix.get(IAgentLifecycleService); + await svc.create({ agentId: 'main' }); + const promptService = svc.handleOf('main')!.accessor.get(IAgentPromptService) as unknown as { + list: ReturnType; + }; + let phase: 'launching' | 'active' | 'empty' = 'launching'; + promptService.list = vi.fn(() => { + if (phase === 'launching') return { launching: true, active: undefined, pending: [] }; + if (phase === 'active') return { launching: false, active: { id: 'p1' }, pending: [] }; + return { launching: false, active: undefined, pending: [] }; + }); + const dispatcher = svc.handleOf('main')!.accessor.get(IEventDispatcher); + const flush = vi.spyOn(dispatcher, 'flush'); + const removal = svc.remove(svc.get('main')!); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(flush).not.toHaveBeenCalled(); + phase = 'active'; + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(flush).not.toHaveBeenCalled(); + phase = 'empty'; + await removal; + expect(flush).toHaveBeenCalled(); + expect(promptDrain.mock.calls.length).toBeGreaterThanOrEqual(2); + }); + + it('remove holds a quiescence guard across the flush and disposal', async () => { + const svc = ix.get(IAgentLifecycleService); + await svc.create({ agentId: 'main' }); + const handle = svc.handleOf('main')!; + const loop = handle.accessor.get(IAgentLoopService) as unknown as { + tryAcquireQuiescence: ReturnType; + }; + const guardDispose = vi.fn(); + loop.tryAcquireQuiescence = vi.fn(() => ({ dispose: guardDispose })); + const dispatcher = handle.accessor.get(IEventDispatcher); + const flush = vi.spyOn(dispatcher, 'flush'); + const dispose = vi.spyOn(handle, 'dispose'); + await svc.remove(svc.get('main')!); + expect(flush).toHaveBeenCalled(); + expect(loop.tryAcquireQuiescence).toHaveBeenCalled(); + expect(guardDispose.mock.invocationCallOrder[0]).toBeGreaterThan( + dispose.mock.invocationCallOrder[0]!, + ); + }); + it('remove keeps the lifecycle context active through async scope teardown', async () => { const svc = ix.get(IAgentLifecycleService); const bus = ix.get(ISessionEventBus); From 36eb4759354a71e75de277a764aa5b7dc1df95d2 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:42:34 +0800 Subject: [PATCH 13/15] fix(kimi-code): stop task producers before the print-mode wire flush Task termination bypasses both the prompt queue and the loop quiescence guard: AgentTaskService dispatches TaskTerminated straight to the wire, and disposal would force-stop still-running tasks after the journals were already flushed. Stop every session agent's tasks up front in the quiesce phase (mirroring AgentLifecycleService.remove()), so each task's termination record is dispatched before the flush instead of racing process.exit. keepAliveOnExit tasks stay exempt, matching remove(). --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 39 +++++++++++++------- apps/kimi-code/test/cli/v2-run-print.test.ts | 15 ++++++-- 2 files changed, 37 insertions(+), 17 deletions(-) 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 2132cbc248..ccd3a8bdab 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -927,21 +927,20 @@ function collectSessionAgentHandles( } /** - * Drain every session agent's prompt queue, cancel queued and active turns, - * then wait for the loops to go idle and the prompt queues to empty. A - * termination signal can arrive mid-turn: without this, the turn's - * cancellation and closing records would only be produced by dispose()'s - * asynchronous teardown — after the wire flush, and after process.exit. - * Draining and cancelling are no-ops on idle agents, so the normal - * (completed/failed turn) exit pays nothing here. + * Stop every session agent's task producers, drain prompt queues, cancel + * queued and active turns, then wait for the loops to go idle and the prompt + * queues to empty. A termination signal can arrive mid-turn: without this, + * the turn's cancellation and closing records would only be produced by + * dispose()'s asynchronous teardown — after the wire flush, and after + * process.exit. Stopping and draining are no-ops on idle agents, so the + * normal (completed/failed turn) exit pays nothing here. * * On success returns a release function: every loop is left holding a - * quiescence guard so late producers (a background-task completion or a cron - * fire, both of which enqueue straight into the loop, bypassing the prompt - * queue) cannot start a new turn — and new records — after the prompt - * queues read empty. The caller holds the release across the wire flush and - * disposal. Returns undefined when quiescence could not be reached within - * the caller's bound. + * quiescence guard so late producers (a cron fire enqueues straight into the + * loop, bypassing the prompt queue) cannot start a new turn — and new + * records — after the prompt queues read empty. The caller holds the release + * across the wire flush and disposal. Returns undefined when quiescence + * could not be reached within the caller's bound. */ async function quiesceSessionAgents( session: ISessionScopeHandle, @@ -965,6 +964,20 @@ async function quiesceSessionAgents( return []; } }); + // Task producers bypass the prompt queue and the loop guard below: their + // termination records are dispatched straight to the wire (and disposal + // would force-stop them after the flush). Stop them up front — settling a + // task dispatches its termination record now, so the flush can persist it + // (mirrors AgentLifecycleService.remove()). + await Promise.allSettled( + handles.flatMap((handle) => { + try { + return [handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed')]; + } catch { + return []; + } + }), + ); // Repeat until a full pass finds every queue empty and every loop freezable: // a prompt can surface after one pass already ran — from startNext's launch // window (the prompt left `pending` and is not yet `active`, so drain could 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 c139615a65..3c1c632038 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -187,7 +187,7 @@ function makeFakeHarness() { 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 () => {}) }], @@ -836,6 +836,13 @@ describe('runV2Print', () => { }); 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; }; @@ -908,14 +915,14 @@ describe('runV2Print', () => { // the loop to have flushed if it were not actually waiting on the queue. await new Promise((resolve) => setTimeout(resolve, 30)); expect(promptService.drain).toHaveBeenCalled(); - expect(order).toEqual(['cancel', 'settled']); + expect(order).toEqual(['stop', 'cancel', 'settled']); promptPhase = 'active'; await new Promise((resolve) => setTimeout(resolve, 30)); - expect(order).toEqual(['cancel', 'settled']); + expect(order).toEqual(['stop', 'cancel', 'settled']); promptPhase = 'empty'; await sigintRun; - expect(order).toEqual(['cancel', 'settled', 'flush', 'exit:130']); + expect(order).toEqual(['stop', 'cancel', 'settled', 'flush', 'exit:130']); // Producers stay frozen across the flush and disposal: the guard taken // during quiesce is only released after app.dispose(). expect(loop.tryAcquireQuiescence).toHaveBeenCalled(); From ed9fe386e349de8986f46508c3a4c43d4ebf1cf2 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:51:07 +0800 Subject: [PATCH 14/15] fix(agent-core-v2): do not let suppression failure short-circuit stopAllOnExit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A detached task's terminal-notification suppression failing (e.g. the persist write rejects) used to reject the whole stopAllOnExit before stopAll() ran, so every task stayed active — and callers settling the rejection (print-mode cleanup, agent removal) proceeded as if the tasks were stopped, losing their termination records at exit. Settle each suppression independently, log the failure, and always stop the tasks. keepAliveOnExit tasks remain exempt by design. --- .../src/agent/task/taskService.ts | 11 ++++++++-- .../test/agent/task/taskService.test.ts | 22 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index b823e457f8..c036ba8c07 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/test/agent/task/taskService.test.ts b/packages/agent-core-v2/test/agent/task/taskService.test.ts index 9f0a10508b..c8f6e90bc5 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); From ab7c82302525cbc1417d45d391e271d99bb68fe1 Mon Sep 17 00:00:00 2001 From: 7Sageer <125936732+7Sageer@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:05:08 +0800 Subject: [PATCH 15/15] chore(kimi-code,agent-core-v2): halve print wire-flush tests and trim comments --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 73 ++------- apps/kimi-code/test/cli/v2-run-print.test.ts | 145 +----------------- .../agentLifecycle/agentLifecycle.test.ts | 68 -------- 3 files changed, 18 insertions(+), 268 deletions(-) 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 ccd3a8bdab..72ed2e9443 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -212,26 +212,14 @@ export async function runV2Print( setCrashPhase('shutdown'); try { await restorePermission(); - // A termination signal can arrive mid-turn: drain pending prompts, - // cancel any queued/active turns, and wait for every agent's loops - // and prompt queues to quiesce first, so the turn's cancellation and - // closing records exist before the flush below drains the journals. - // No-op when every agent is already idle. + // A termination signal can arrive mid-turn: cancel turns and wait for idle agents first. await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}); } finally { try { - // The shutdown phases are independent of each other; run them - // concurrently so their individual allowances cannot sum past the - // outer PROMPT_CLEANUP_TIMEOUT_MS bound and let the caller's - // process.exit cut off the tail (app.dispose included). + // Concurrent so the phases' allowances cannot sum past PROMPT_CLEANUP_TIMEOUT_MS. await Promise.all([ - // The turn's tail records (step.end / turn.ended / prompt.completed) - // are dispatched fire-and-forget and reach the journal only through - // the wire service's async persist queue. Without an explicit flush, - // a cleanup that returns fast (e.g. telemetry disabled) lets - // process.exit cut off that queue before the records land on disk. - // Best-effort: a persist failure was already reported where the - // append failed. + // 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) @@ -240,9 +228,7 @@ export async function runV2Print( ]); app.dispose(); } finally { - // Quiescence guards stay held across the flush and disposal so a - // late background-task completion or cron fire cannot start a new - // turn (and new records) after the journals were drained. + // Keep producers frozen until the journals are drained and disposed. releaseQuiescence?.(); } } @@ -907,12 +893,7 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number { return count; } -/** - * Every agent handle in the session. The main agent handle is included - * explicitly: `IAgentLifecycleService` skips `closing` agents, but a closing - * agent's in-flight turn and already-dispatched tail records still deserve to - * land on disk. - */ +/** 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, @@ -927,20 +908,8 @@ function collectSessionAgentHandles( } /** - * Stop every session agent's task producers, drain prompt queues, cancel - * queued and active turns, then wait for the loops to go idle and the prompt - * queues to empty. A termination signal can arrive mid-turn: without this, - * the turn's cancellation and closing records would only be produced by - * dispose()'s asynchronous teardown — after the wire flush, and after - * process.exit. Stopping and draining are no-ops on idle agents, so the - * normal (completed/failed turn) exit pays nothing here. - * - * On success returns a release function: every loop is left holding a - * quiescence guard so late producers (a cron fire enqueues straight into the - * loop, bypassing the prompt queue) cannot start a new turn — and new - * records — after the prompt queues read empty. The caller holds the release - * across the wire flush and disposal. Returns undefined when quiescence - * could not be reached within the caller's bound. + * 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, @@ -959,16 +928,12 @@ async function quiesceSessionAgents( try { return [handle.accessor.get(IAgentLoopService)]; } catch { - // A torn-down agent scope has no loop to quiesce; the wire flush below - // still covers its already-dispatched records. + // A torn-down agent scope has no loop to quiesce. return []; } }); - // Task producers bypass the prompt queue and the loop guard below: their - // termination records are dispatched straight to the wire (and disposal - // would force-stop them after the flush). Stop them up front — settling a - // task dispatches its termination record now, so the flush can persist it - // (mirrors AgentLifecycleService.remove()). + // 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 { @@ -978,14 +943,8 @@ async function quiesceSessionAgents( } }), ); - // Repeat until a full pass finds every queue empty and every loop freezable: - // a prompt can surface after one pass already ran — from startNext's launch - // window (the prompt left `pending` and is not yet `active`, so drain could - // not see it) or from a cancelled turn's settle chain. The snapshot covers - // all three states; the loop reports idle (releaseActiveTurn) before the - // prompt-settle chain dispatches the final record, and settle() clears the - // active prompt and dispatches that record in one synchronous block, so an - // empty snapshot proves it was already queued for the wire flush. + // 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) { @@ -1033,11 +992,7 @@ async function quiesceSessionAgents( } } -/** - * Flush every session agent's wire journal. Each flush settles independently: - * one agent's broken journal must not cut the wait short for the others (the - * caller proceeds to `process.exit`). - */ +/** Flush every session agent's wire journal; each flush settles independently. */ async function flushSessionWires( session: ISessionScopeHandle, mainAgent: IAgentScopeHandle, 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 3c1c632038..e3b12c1b70 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -713,107 +713,6 @@ describe('runV2Print', () => { expect(app.dispose).toHaveBeenCalled(); }); - it('keeps waiting for healthy agents when another agent\'s flush fails', async () => { - const stdout = writer(); - const stderr = writer(); - const { app, agent, agentServices, sessionServices } = 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 order: string[] = []; - const mainDispatcher = agentServices.get(IEventDispatcher) as { - flush: ReturnType; - }; - mainDispatcher.flush.mockRejectedValueOnce(new Error('disk full')); - const subAgent = fakeScope( - 'sub', - new Map([ - [ - IEventDispatcher, - { - flush: vi.fn(async () => { - await new Promise((resolve) => setTimeout(resolve, 10)); - order.push('sub-flushed'); - }), - }, - ], - ]), - ); - const lifecycle = sessionServices.get(IAgentLifecycleService) as { - list: ReturnType; - handleOf: ReturnType; - }; - lifecycle.list.mockReturnValue([{ agentId: 'main' }, { agentId: 'sub' }]); - lifecycle.handleOf.mockImplementation((id: string) => (id === 'sub' ? subAgent : agent)); - app.dispose.mockImplementation(() => { - order.push('disposed'); - }); - - 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(order).toEqual(['sub-flushed', 'disposed']); - }); - - it('runs wire flush and telemetry shutdown concurrently', async () => { - const stdout = writer(); - const stderr = writer(); - const { app, agentServices, appServices } = makeFakeHarness(); - - const dispatcher = agentServices.get(IEventDispatcher) as { - flush: ReturnType; - }; - let releaseFlush!: () => void; - const flushGate = new Promise((resolve) => { - releaseFlush = resolve; - }); - dispatcher.flush.mockReturnValueOnce(flushGate); - const telemetry = appServices.get(ITelemetryService) as { - shutdown: ReturnType; - }; - let releaseTelemetry!: () => void; - const telemetryGate = new Promise((resolve) => { - releaseTelemetry = resolve; - }); - telemetry.shutdown.mockReturnValueOnce(telemetryGate); - - mocks.bootstrap.mockReturnValue({ app }); - mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); - - const run = runV2Print(opts() as never, '1.2.3-test', { stdout, stderr }); - // A sequential cleanup would sit in the wire flush's 3s allowance before - // even starting telemetry shutdown; concurrent phases are both in flight - // well within that window. - for ( - let i = 0; - i < 200 && - (dispatcher.flush.mock.calls.length === 0 || telemetry.shutdown.mock.calls.length === 0); - i++ - ) { - await new Promise((resolve) => setTimeout(resolve, 10)); - } - expect(dispatcher.flush).toHaveBeenCalled(); - expect(telemetry.shutdown).toHaveBeenCalled(); - releaseFlush(); - releaseTelemetry(); - await run; - expect(app.dispose).toHaveBeenCalled(); - }); - it('cancels and settles the active turn before flushing on a termination signal', async () => { const stdout = writer(); const stderr = writer(); @@ -850,10 +749,8 @@ describe('runV2Print', () => { order.push('flush'); }); - // A turn that is still in flight when the signal arrives. The prompt - // queue reports the launch window first (the prompt left `pending` and is - // not yet `active`), then the running prompt, and only goes empty once - // released below. + // 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; @@ -905,14 +802,10 @@ describe('runV2Print', () => { const onSigint = handlers.get('SIGINT')!; settleTurn({ type: 'cancelled', steps: 0, reason: new Error('aborted') }); const sigintRun = onSigint(); - // Quiesce waits for the prompt queue to empty before the wire flush: the - // loops are already idle, but while the snapshot still reports the launch - // window or an active prompt the flush must not happen. + // 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)); } - // A few quiesce poll cycles (the poll interval is 10ms): enough time for - // the loop to have flushed if it were not actually waiting on the queue. await new Promise((resolve) => setTimeout(resolve, 30)); expect(promptService.drain).toHaveBeenCalled(); expect(order).toEqual(['stop', 'cancel', 'settled']); @@ -923,41 +816,11 @@ describe('runV2Print', () => { await sigintRun; expect(order).toEqual(['stop', 'cancel', 'settled', 'flush', 'exit:130']); - // Producers stay frozen across the flush and disposal: the guard taken - // during quiesce is only released after app.dispose(). + // 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); }); - - it('retries quiescence acquisition while a producer keeps the loop busy', async () => { - const stdout = writer(); - const stderr = writer(); - const { app, agentServices } = makeFakeHarness(); - - const loop = agentServices.get(IAgentLoopService) as { - tryAcquireQuiescence: ReturnType; - }; - let attempts = 0; - loop.tryAcquireQuiescence = vi.fn(() => { - attempts += 1; - // The first pass finds the loop busy (a background-task completion just - // enqueued a turn); the second pass can freeze it. - return attempts === 1 ? undefined : { dispose: vi.fn() }; - }); - - mocks.bootstrap.mockReturnValue({ app }); - mocks.ensureMainAgent.mockResolvedValue({ agentId: 'main', generation: 1 }); - - await runV2Print(opts() as never, '1.2.3-test', { stdout, stderr }); - - expect(attempts).toBeGreaterThanOrEqual(2); - const dispatcher = agentServices.get(IEventDispatcher) as { - flush: ReturnType; - }; - expect(dispatcher.flush).toHaveBeenCalled(); - expect(app.dispose).toHaveBeenCalled(); - }); }); 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 6567041425..b1c6c05a38 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -520,74 +520,6 @@ describe('AgentLifecycleService', () => { expect(flush).toHaveBeenCalled(); }); - it('remove waits for the prompt queue to go idle before flushing', async () => { - const svc = ix.get(IAgentLifecycleService); - await svc.create({ agentId: 'main' }); - const promptService = svc.handleOf('main')!.accessor.get(IAgentPromptService) as unknown as { - list: ReturnType; - }; - let queueIdle = false; - promptService.list = vi.fn(() => - queueIdle - ? { launching: false, active: undefined, pending: [] } - : { launching: false, active: { id: 'p1' }, pending: [] }, - ); - const dispatcher = svc.handleOf('main')!.accessor.get(IEventDispatcher); - const flush = vi.spyOn(dispatcher, 'flush'); - const removal = svc.remove(svc.get('main')!); - await new Promise((resolve) => setTimeout(resolve, 50)); - expect(flush).not.toHaveBeenCalled(); - queueIdle = true; - await removal; - expect(flush).toHaveBeenCalled(); - }); - - it('remove re-cancels a prompt that finishes launching during the removal', async () => { - const svc = ix.get(IAgentLifecycleService); - await svc.create({ agentId: 'main' }); - const promptService = svc.handleOf('main')!.accessor.get(IAgentPromptService) as unknown as { - list: ReturnType; - }; - let phase: 'launching' | 'active' | 'empty' = 'launching'; - promptService.list = vi.fn(() => { - if (phase === 'launching') return { launching: true, active: undefined, pending: [] }; - if (phase === 'active') return { launching: false, active: { id: 'p1' }, pending: [] }; - return { launching: false, active: undefined, pending: [] }; - }); - const dispatcher = svc.handleOf('main')!.accessor.get(IEventDispatcher); - const flush = vi.spyOn(dispatcher, 'flush'); - const removal = svc.remove(svc.get('main')!); - await new Promise((resolve) => setTimeout(resolve, 50)); - expect(flush).not.toHaveBeenCalled(); - phase = 'active'; - await new Promise((resolve) => setTimeout(resolve, 50)); - expect(flush).not.toHaveBeenCalled(); - phase = 'empty'; - await removal; - expect(flush).toHaveBeenCalled(); - expect(promptDrain.mock.calls.length).toBeGreaterThanOrEqual(2); - }); - - it('remove holds a quiescence guard across the flush and disposal', async () => { - const svc = ix.get(IAgentLifecycleService); - await svc.create({ agentId: 'main' }); - const handle = svc.handleOf('main')!; - const loop = handle.accessor.get(IAgentLoopService) as unknown as { - tryAcquireQuiescence: ReturnType; - }; - const guardDispose = vi.fn(); - loop.tryAcquireQuiescence = vi.fn(() => ({ dispose: guardDispose })); - const dispatcher = handle.accessor.get(IEventDispatcher); - const flush = vi.spyOn(dispatcher, 'flush'); - const dispose = vi.spyOn(handle, 'dispose'); - await svc.remove(svc.get('main')!); - expect(flush).toHaveBeenCalled(); - expect(loop.tryAcquireQuiescence).toHaveBeenCalled(); - expect(guardDispose.mock.invocationCallOrder[0]).toBeGreaterThan( - dispose.mock.invocationCallOrder[0]!, - ); - }); - it('remove keeps the lifecycle context active through async scope teardown', async () => { const svc = ix.get(IAgentLifecycleService); const bus = ix.get(ISessionEventBus);