fix(cli): preserve buffered prompt output during signal shutdown

.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.

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<void> => {};
   let removeTerminationCleanup: (() => void) | undefined;
   let cleanupPromise: Promise<void> | undefined;
+  let outputDrainPromise: Promise<void> | undefined;
   const cleanupPromptRun = async (): Promise<void> => {
     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<void>,

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<void> => {};
   let removeTerminationCleanup: (() => void) | undefined;
   let cleanupPromise: Promise<void> | undefined;
+  let outputDrainPromise: Promise<void> | undefined;
   let telemetryService: ITelemetryService | undefined;
   const cleanup = async (): Promise<void> => {
     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);
 

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<void>((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<typeof runPrompt>[2] & { process: ReturnType<typeof fakeProcess> });
+
+    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;