diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index bfb3c0b90..b587daff0 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/apps/kimi-code/src/cli/run-prompt.ts @@ -236,6 +236,7 @@ export async function runPrompt( resolveKeepAliveOnExit(config, process.env), resolvePrintWaitCeilingMs(config), outputFormat, + finalOnly, stdout, stderr, io.clock ?? defaultPromptClock, @@ -1170,6 +1171,7 @@ async function drainBackgroundTasksOnExit( keepAliveOnExit: boolean, ceilingMs: number, outputFormat: PromptOutputFormat, + finalOnly: boolean, stdout: PromptOutput, stderr: PromptOutput, clock: PromptClock, @@ -1177,8 +1179,10 @@ async function drainBackgroundTasksOnExit( if (keepAliveOnExit) return; let active = await session.listBackgroundTasks({ activeOnly: true }); if (active.length === 0) return; + // `--final-message-only` promises a single final assistant message per turn, + // so the drain still waits for tasks but does not emit their notifications. const unsubscribe = - outputFormat === 'stream-json' + outputFormat === 'stream-json' && !finalOnly ? session.onEvent((event) => { const notification = toNotificationMessage(event); if (notification?.event === 'background.task.terminated') { diff --git a/apps/kimi-code/test/cli/run-prompt.test.ts b/apps/kimi-code/test/cli/run-prompt.test.ts index 494742edc..ff002e084 100644 --- a/apps/kimi-code/test/cli/run-prompt.test.ts +++ b/apps/kimi-code/test/cli/run-prompt.test.ts @@ -1399,6 +1399,53 @@ describe('runPrompt', () => { ); }); + it('waits for background tasks but suppresses drain notifications in final-message-only mode', async () => { + let listCalls = 0; + mocks.session.listBackgroundTasks.mockImplementation(async () => { + listCalls += 1; + return listCalls === 1 ? [{ taskId: 'b1', status: 'running' }] : []; + }); + const fakeClock = { + now: () => 0, + sleep: vi.fn(async () => { + for (const handler of Array.from(mocks.eventHandlers)) { + handler( + mocks.mainEvent({ + type: 'background.task.terminated', + info: { + taskId: 'b1', + kind: 'process', + status: 'completed', + description: 'build', + startedAt: 0, + endedAt: 1, + }, + }), + ); + } + }), + }; + mocks.session.prompt.mockImplementationOnce(async () => { + for (const handler of mocks.eventHandlers) { + handler(mocks.mainEvent({ type: 'turn.started', turnId: 80, origin: { kind: 'user' } })); + handler(mocks.mainEvent({ type: 'assistant.delta', turnId: 80, delta: 'done' })); + handler(mocks.mainEvent({ type: 'turn.ended', turnId: 80, reason: 'completed' })); + } + }); + const stdout = writer(); + + await runPrompt(opts({ outputFormat: 'stream-json', finalMessageOnly: true }), '1.2.3-test', { + stdout, + stderr: writer(), + clock: fakeClock, + }); + + // Still waits for the task, but emits nothing beyond the single final message. + expect(fakeClock.sleep).toHaveBeenCalled(); + expect(stdout.text()).toBe('{"role":"assistant","content":"done"}\n'); + expect(stdout.text()).not.toContain('notification'); + }); + it('skips the background-task wait when keepAliveOnExit is set via env (Tier 2)', async () => { process.env['KIMI_CODE_BACKGROUND_KEEP_ALIVE_ON_EXIT'] = '1'; mocks.session.listBackgroundTasks.mockResolvedValue([{ taskId: 'b1', status: 'running' }]);