diff --git a/.changeset/fix-headless-signal-output.md b/.changeset/fix-headless-signal-output.md new file mode 100644 index 0000000000..35be4e0704 --- /dev/null +++ b/.changeset/fix-headless-signal-output.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Preserve buffered headless output when a signal arrives during prompt cleanup. diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index abee299622..de09db1399 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/apps/kimi-code/src/cli/run-prompt.ts @@ -14,11 +14,13 @@ import { type SessionStatus, type TelemetryClient, } from '@moonshot-ai/kimi-code-sdk'; +import { Writable } from 'node:stream'; import { resolve } from 'pathe'; import { CLI_SHUTDOWN_TIMEOUT_MS, PROMPT_CLEANUP_TIMEOUT_MS } from '#/constant/app'; import { isKimiV2Enabled } from './experimental-v2'; +import { drainStdio } from './headless-exit'; import { resolveOutputFormat } from './options'; import type { CLIOptions, PromptOutputFormat } from './options'; import { @@ -113,6 +115,9 @@ export async function runPrompt( const stdout = io.stdout ?? process.stdout; const stderr = io.stderr ?? process.stderr; const promptProcess = io.process ?? process; + const streams = [stdout, stderr].filter( + (stream): stream is Writable => stream instanceof Writable, + ); const outputFormat = resolveOutputFormat(opts); const workDir = process.cwd(); const telemetryBootstrap = createCliTelemetryBootstrap(); @@ -146,9 +151,9 @@ export async function runPrompt( let restorePromptSessionPermission = async (): Promise => {}; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; + let outputDrainPromise: Promise | undefined; const cleanupPromptRun = async (): Promise => { const pending = (cleanupPromise ??= (async () => { - removeTerminationCleanup?.(); setCrashPhase('shutdown'); try { await restorePromptSessionPermission(); @@ -162,9 +167,24 @@ export async function runPrompt( // keep a completed headless run alive forever. The cleanup keeps running in // the background if it overruns; the caller (`kimi -p`) force-exits shortly // after, so any straggling work is torn down with the process. - await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS); + const cleanupOutcome = await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS).then( + () => ({ ok: true }) as const, + (error: unknown) => ({ ok: false, error }) as const, + ); + await (outputDrainPromise ??= drainStdio(streams).finally(() => { + // Keep the signal handlers installed through async cleanup and output + // flushing. Removing them earlier restores Node's default immediate + // signal exit, which can discard the final assistant and resume hint. + removeTerminationCleanup?.(); + })); + if (!cleanupOutcome.ok) { + throw cleanupOutcome.error; + } }; - removeTerminationCleanup = installPromptTerminationCleanup(promptProcess, cleanupPromptRun); + removeTerminationCleanup = installPromptTerminationCleanup( + promptProcess, + cleanupPromptRun, + ); try { await harness.ensureConfigFile(); @@ -422,6 +442,10 @@ function installHeadlessHandlers(session: PromptSession): void { session.setQuestionHandler(() => null); } +/** + * Install one-shot signal handlers that clean up and flush queued output before + * forcing the process to exit with the signal's conventional exit code. + */ export function installPromptTerminationCleanup( promptProcess: PromptProcess, cleanup: () => Promise, 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 ff0c96c590..a78ace8606 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -17,6 +17,7 @@ */ import { readFile } from 'node:fs/promises'; +import { Writable } from 'node:stream'; import { IAgentGoalService, @@ -73,6 +74,7 @@ import { parseHeadlessGoalCreate, type HeadlessGoalCreate, } from '../goal-prompt'; +import { drainStdio } from '../headless-exit'; import { type PromptRunIO, configuredModel, @@ -112,6 +114,9 @@ export async function runV2Print( const stdout = io.stdout ?? process.stdout; const stderr = io.stderr ?? process.stderr; const promptProcess = io.process ?? process; + const streams = [stdout, stderr].filter( + (stream): stream is Writable => stream instanceof Writable, + ); const outputFormat = resolveOutputFormat(opts); const workDir = process.cwd(); @@ -164,10 +169,10 @@ export async function runV2Print( let restorePermission = async (): Promise => {}; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; + let outputDrainPromise: Promise | undefined; let telemetryService: ITelemetryService | undefined; const cleanup = async (): Promise => { const pending = (cleanupPromise ??= (async () => { - removeTerminationCleanup?.(); try { await restorePermission(); } finally { @@ -177,7 +182,19 @@ export async function runV2Print( app.dispose(); } })()); - await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS); + const cleanupOutcome = await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS).then( + () => ({ ok: true }) as const, + (error: unknown) => ({ ok: false, error }) as const, + ); + await (outputDrainPromise ??= drainStdio(streams).finally(() => { + // Keep the signal handlers installed through async cleanup and output + // flushing. Removing them earlier restores Node's default immediate + // signal exit, which can discard the final assistant and resume hint. + removeTerminationCleanup?.(); + })); + if (!cleanupOutcome.ok) { + throw cleanupOutcome.error; + } }; removeTerminationCleanup = installPromptTerminationCleanup(promptProcess, cleanup); diff --git a/apps/kimi-code/test/cli/run-prompt.test.ts b/apps/kimi-code/test/cli/run-prompt.test.ts index 4bad127d89..9b2d398376 100644 --- a/apps/kimi-code/test/cli/run-prompt.test.ts +++ b/apps/kimi-code/test/cli/run-prompt.test.ts @@ -1,4 +1,5 @@ import type { createKimiDeviceId as createKimiDeviceIdFn } from '@moonshot-ai/kimi-code-oauth'; +import { Writable } from 'node:stream'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { runPrompt } from '#/cli/run-prompt'; @@ -1033,6 +1034,106 @@ describe('runPrompt', () => { expect(mocks.harnessClose).toHaveBeenCalledTimes(1); }); + it('preserves completed stream-json output when signalled during cleanup', async () => { + let releaseCleanup: (() => void) | undefined; + mocks.harnessClose.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseCleanup = resolve; + }), + ); + + const chunks: string[] = []; + let blocked = true; + let releaseBufferedWrite: (() => void) | undefined; + const stdout = new Writable({ + highWaterMark: 1, + write(chunk, _encoding, callback) { + chunks.push(chunk.toString()); + if (!blocked) { + callback(); + return; + } + releaseBufferedWrite = callback; + }, + }); + + const processMock = fakeProcess(); + const run = runPrompt(opts({ outputFormat: 'stream-json' }), '1.2.3-test', { + stdout, + stderr: writer(), + process: processMock, + } as Parameters[2] & { process: ReturnType }); + + await waitForAssertion(() => { + expect(mocks.harnessClose).toHaveBeenCalled(); + }); + expect(stdout.writableNeedDrain).toBe(true); + const signalHandler = processMock.listener('SIGTERM'); + if (signalHandler === undefined) { + blocked = false; + releaseBufferedWrite?.(); + releaseCleanup?.(); + await run; + } + expect(signalHandler).toBeDefined(); + + const termination = signalHandler?.(); + releaseCleanup?.(); + await Promise.resolve(); + expect(processMock.exit).not.toHaveBeenCalled(); + + blocked = false; + releaseBufferedWrite?.(); + await termination; + await run; + + expect(chunks.join('')).toContain('{"role":"assistant","content":"hello world"}\n'); + expect(chunks.join('')).toContain('"type":"session.resume_hint"'); + expect(processMock.exit).toHaveBeenCalledWith(143); + }); + + it('propagates cleanup failures after buffered output finishes draining', async () => { + vi.useFakeTimers(); + try { + mocks.harnessClose.mockRejectedValueOnce(new Error('close failed')); + + let releaseBufferedWrite: (() => void) | undefined; + let blocked = true; + const stdout = new Writable({ + highWaterMark: 1, + write(_chunk, _encoding, callback) { + if (!blocked) { + callback(); + return; + } + releaseBufferedWrite = callback; + }, + }); + const run = runPrompt(opts({ outputFormat: 'stream-json' }), '1.2.3-test', { + stdout, + stderr: writer(), + }); + + for ( + let attempt = 0; + attempt < 20 && mocks.harnessClose.mock.calls.length === 0; + attempt += 1 + ) { + await Promise.resolve(); + } + expect(mocks.harnessClose).toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(PROMPT_CLEANUP_TIMEOUT_MS); + blocked = false; + releaseBufferedWrite?.(); + + await expect(run).rejects.toThrow('close failed'); + } finally { + vi.useRealTimers(); + } + }); + it('waits for the pending auto permission write before signal restore', async () => { let releaseAutoPermission!: () => void; let releasePrompt!: () => void;