From fbc128b6821fae0566fe90d0121eb65b91542515 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Sat, 25 Jul 2026 02:35:25 +0000 Subject: [PATCH 1/2] fix(cli): preserve buffered output during signal shutdown --- .changeset/fix-headless-signal-output.md | 5 + apps/kimi-code/src/cli/run-prompt.ts | 33 ++++++- apps/kimi-code/src/cli/v2/run-v2-print.ts | 23 ++++- apps/kimi-code/test/cli/run-prompt.test.ts | 101 +++++++++++++++++++++ 4 files changed, 156 insertions(+), 6 deletions(-) create mode 100644 .changeset/fix-headless-signal-output.md 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..f9b554deed 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,25 @@ 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, + streams, + ); try { await harness.ensureConfigFile(); @@ -422,9 +443,14 @@ 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, + streams: readonly Writable[], ): () => void { let terminating = false; const exitAfterCleanup = async (signal: NodeJS.Signals): Promise => { @@ -433,6 +459,7 @@ export function installPromptTerminationCleanup( try { await cleanup(); } finally { + await drainStdio(streams); promptProcess.exit(signalExitCode(signal)); } }; diff --git a/apps/kimi-code/src/cli/v2/run-v2-print.ts b/apps/kimi-code/src/cli/v2/run-v2-print.ts index ff0c96c590..48cf74dd07 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,9 +182,21 @@ 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); + removeTerminationCleanup = installPromptTerminationCleanup(promptProcess, cleanup, streams); try { // Install the appender BEFORE resolving the session: `session_started` and 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; From 1fcd71174605dea398ade5f6ce7e9e39d42af6d2 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Sat, 25 Jul 2026 02:53:51 +0000 Subject: [PATCH 2/2] fix(cli): drain signal output once --- apps/kimi-code/src/cli/run-prompt.ts | 3 --- apps/kimi-code/src/cli/v2/run-v2-print.ts | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index f9b554deed..de09db1399 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/apps/kimi-code/src/cli/run-prompt.ts @@ -184,7 +184,6 @@ export async function runPrompt( removeTerminationCleanup = installPromptTerminationCleanup( promptProcess, cleanupPromptRun, - streams, ); try { @@ -450,7 +449,6 @@ function installHeadlessHandlers(session: PromptSession): void { export function installPromptTerminationCleanup( promptProcess: PromptProcess, cleanup: () => Promise, - streams: readonly Writable[], ): () => void { let terminating = false; const exitAfterCleanup = async (signal: NodeJS.Signals): Promise => { @@ -459,7 +457,6 @@ export function installPromptTerminationCleanup( try { await cleanup(); } finally { - await drainStdio(streams); promptProcess.exit(signalExitCode(signal)); } }; diff --git a/apps/kimi-code/src/cli/v2/run-v2-print.ts b/apps/kimi-code/src/cli/v2/run-v2-print.ts index 48cf74dd07..a78ace8606 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -196,7 +196,7 @@ export async function runV2Print( throw cleanupOutcome.error; } }; - removeTerminationCleanup = installPromptTerminationCleanup(promptProcess, cleanup, streams); + removeTerminationCleanup = installPromptTerminationCleanup(promptProcess, cleanup); try { // Install the appender BEFORE resolving the session: `session_started` and