fix(cli): force-exit headless runs so a lingering handle can't wedge kimi -p

.changeset/fix-headless-force-exit.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Force-exit headless runs (`kimi -p`) so a stray ref'd handle left over from the run can't keep a completed run alive until an external timeout, and bound prompt cleanup so a wedged shutdown step can't hang shutdown.

apps/kimi-code/src/cli/headless-exit.ts

@@ -0,0 +1,96 @@
+import type { Writable } from 'node:stream';
+
+import { HEADLESS_FORCE_EXIT_GRACE_MS, HEADLESS_STDIO_DRAIN_TIMEOUT_MS } from '#/constant/app';
+
+/** Minimal process surface needed to force a headless run to terminate. */
+export interface ExitableProcess {
+  exit(code?: number): void;
+}
+
+/**
+ * Schedule a best-effort force-exit for a completed headless (`kimi -p`) run.
+ *
+ * Print mode does not call `process.exit()`; it relies on the Node event loop
+ * draining once the run is done. If a stray ref'd handle survives shutdown — a
+ * lingering socket (e.g. a connection blackholed by a restrictive firewall, or
+ * an HTTP/2 session kept alive by PING), an un-cleared timer, or a child whose
+ * pipes stay open — the loop never empties and the process hangs until an
+ * external timeout kills it.
+ *
+ * This arms an **unref'd** fallback timer: a healthy run drains and exits
+ * naturally before it fires (so behaviour is unchanged), and the timer itself
+ * never keeps the loop alive. It only force-exits a run whose loop is already
+ * wedged. The exit code is read lazily at fire time so callers may set
+ * `process.exitCode` after scheduling (e.g. a goal turn mapping its terminal
+ * status to a non-zero code).
+ *
+ * Returns the timer handle so callers/tests can `clearTimeout` it.
+ */
+export function scheduleHeadlessForceExit(
+  proc: ExitableProcess,
+  getExitCode: () => number,
+  graceMs: number = HEADLESS_FORCE_EXIT_GRACE_MS,
+): NodeJS.Timeout {
+  const timer = setTimeout(() => {
+    proc.exit(getExitCode());
+  }, graceMs);
+  timer.unref?.();
+  return timer;
+}
+
+/** Resolve once a stream's currently-buffered writes have flushed to its sink. */
+function flushStream(stream: Writable): Promise<void> {
+  return new Promise<void>((resolve) => {
+    try {
+      // An empty write's callback fires after all previously-queued writes have
+      // been flushed (writes are ordered), which is the documented way to know a
+      // stream's buffer has drained.
+      stream.write('', () => resolve());
+    } catch {
+      resolve();
+    }
+  });
+}
+
+/**
+ * Wait for buffered output on the given streams to flush, bounded by `timeoutMs`.
+ *
+ * A slow or piped consumer that hasn't read all of stdout/stderr yet leaves the
+ * pipe as a legitimate ref'd handle keeping the loop alive. Flushing before any
+ * force-exit prevents truncating output from an otherwise-successful run. The
+ * wait is bounded so a permanently-stuck consumer can't re-introduce the hang.
+ */
+export async function drainStdio(
+  streams: readonly Writable[],
+  timeoutMs: number = HEADLESS_STDIO_DRAIN_TIMEOUT_MS,
+): Promise<void> {
+  let timer: NodeJS.Timeout | undefined;
+  const timeout = new Promise<void>((resolve) => {
+    timer = setTimeout(resolve, timeoutMs);
+    timer.unref?.();
+  });
+  try {
+    await Promise.race([Promise.all(streams.map(flushStream)).then(() => undefined), timeout]);
+  } finally {
+    if (timer !== undefined) clearTimeout(timer);
+  }
+}
+
+/**
+ * Finalize a completed headless run: flush stdio, then arm the force-exit
+ * backstop.
+ *
+ * Draining first means in-flight legitimate output is fully written before the
+ * backstop can fire, and — since drained stdio no longer holds the loop — only a
+ * genuinely leaked handle can keep it alive afterwards, which is exactly what
+ * the backstop is for.
+ */
+export async function finalizeHeadlessRun(
+  proc: ExitableProcess,
+  streams: readonly Writable[],
+  getExitCode: () => number,
+  options: { drainTimeoutMs?: number; graceMs?: number } = {},
+): Promise<void> {
+  await drainStdio(streams, options.drainTimeoutMs ?? HEADLESS_STDIO_DRAIN_TIMEOUT_MS);
+  scheduleHeadlessForceExit(proc, getExitCode, options.graceMs);
+}

apps/kimi-code/src/cli/run-prompt.ts

@@ -19,7 +19,7 @@ import {
 } from '@moonshot-ai/kimi-code-sdk';
 import { resolve } from 'pathe';
 
-import { CLI_SHUTDOWN_TIMEOUT_MS } from '#/constant/app';
+import { CLI_SHUTDOWN_TIMEOUT_MS, PROMPT_CLEANUP_TIMEOUT_MS } from '#/constant/app';
 
 import type { CLIOptions, PromptOutputFormat } from './options';
 import {
@@ -32,6 +32,44 @@ import {
 import { createCliTelemetryBootstrap, initializeCliTelemetry } from './telemetry';
 import { createKimiCodeHostIdentity } from './version';
 
+/**
+ * Await `promise`, but stop waiting after `timeoutMs`.
+ *
+ * The timeout only bounds how long we WAIT — it does not change the outcome:
+ *  - if `promise` settles first, its result is propagated (a rejection throws),
+ *    so a cleanup step that actually fails in time still surfaces;
+ *  - if the timeout wins, we resolve (give up waiting) and swallow the abandoned
+ *    promise's eventual late rejection so it can't surface as an unhandled
+ *    rejection.
+ *
+ * Used to bound shutdown so a wedged cleanup step can't keep a completed
+ * headless run alive, without silently swallowing a cleanup that fails fast. The
+ * timer is unref'd so it never keeps the loop alive on its own.
+ */
+async function raceWithTimeout(promise: Promise<void>, timeoutMs: number): Promise<void> {
+  let timedOut = false;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  // Attach the catch eagerly (synchronously) so `promise` is always consumed and
+  // a late rejection can never become an unhandled rejection. Before the timeout
+  // wins, the handler rethrows so a real cleanup failure still propagates.
+  const guarded = promise.catch((error: unknown) => {
+    if (timedOut) return;
+    throw error;
+  });
+  const timedOutSignal = new Promise<void>((resolve) => {
+    timer = setTimeout(() => {
+      timedOut = true;
+      resolve();
+    }, timeoutMs);
+    timer.unref?.();
+  });
+  try {
+    await Promise.race([guarded, timedOutSignal]);
+  } finally {
+    if (timer !== undefined) clearTimeout(timer);
+  }
+}
+
 interface PromptOutput {
   readonly columns?: number | undefined;
   write(chunk: string): boolean;
@@ -96,7 +134,7 @@ export async function runPrompt(
   let removeTerminationCleanup: (() => void) | undefined;
   let cleanupPromise: Promise<void> | undefined;
   const cleanupPromptRun = async (): Promise<void> => {
-    cleanupPromise ??= (async () => {
+    const pending = (cleanupPromise ??= (async () => {
       removeTerminationCleanup?.();
       setCrashPhase('shutdown');
       try {
@@ -105,8 +143,13 @@ export async function runPrompt(
         await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS });
         await harness.close();
       }
-    })();
-    await cleanupPromise;
+    })());
+    // Bound cleanup so a wedged shutdown step (e.g. a SessionEnd hook, MCP
+    // shutdown, or a connection blackholed by a restrictive firewall) cannot
+    // 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);
   };
   removeTerminationCleanup = installPromptTerminationCleanup(promptProcess, cleanupPromptRun);
 

apps/kimi-code/src/constant/app.ts

@@ -14,6 +14,26 @@ export const WEB_UI_MODE = 'web';
 // Give telemetry a short flush window without making CLI exit feel stuck.
 export const CLI_SHUTDOWN_TIMEOUT_MS = 3000;
 
+// Upper bound on headless (`kimi -p`) shutdown. A wedged cleanup step (e.g. a
+// SessionEnd hook, an MCP shutdown, or a connection blackholed by a restrictive
+// firewall) must not keep a completed run alive indefinitely — once this elapses
+// we stop waiting on cleanup and let the run return.
+export const PROMPT_CLEANUP_TIMEOUT_MS = 8000;
+
+// Grace after a headless run has fully completed (turn done, cleanup attempted)
+// before force-exiting. `kimi -p` otherwise relies on the event loop draining to
+// exit; a stray ref'd handle (socket/timer/child) left over from the run would
+// wedge it. The guard timer is unref'd, so a healthy run still exits naturally
+// well before this fires.
+export const HEADLESS_FORCE_EXIT_GRACE_MS = 2000;
+
+// Max time to wait for buffered stdout/stderr to flush before arming the
+// force-exit fallback. A slow/piped consumer's still-draining stdio is a
+// legitimate ref'd handle — flushing first prevents the fallback from
+// truncating completed output. Bounded so a permanently-stuck consumer can't
+// re-introduce the hang.
+export const HEADLESS_STDIO_DRAIN_TIMEOUT_MS = 10000;
+
 // Published npm package name; this can differ from the executable command.
 export const NPM_PACKAGE_NAME = '@moonshot-ai/kimi-code';
 

apps/kimi-code/src/main.ts

@@ -23,6 +23,7 @@ import {
 } from '@moonshot-ai/kimi-telemetry';
 
 import { createProgram } from './cli/commands';
+import { finalizeHeadlessRun } from './cli/headless-exit';
 import type { CLIOptions } from './cli/options';
 import { OptionConflictError, validateOptions } from './cli/options';
 import { runPrompt } from './cli/run-prompt';
@@ -38,7 +39,22 @@ import { cleanupStaleNativeCacheForCurrent } from './native/native-assets';
 import { installNativeModuleHook } from './native/module-hook';
 import { runNativeAssetSmokeIfRequested } from './native/smoke';
 
-export async function handleMainCommand(opts: CLIOptions, version: string): Promise<void> {
+/**
+ * Outcome of a CLI command run, reported back to the process entrypoint.
+ *
+ * `handleMainCommand` is a reusable, unit-tested handler — it must not terminate
+ * the process itself. It reports here whether a headless (`kimi -p`) run
+ * completed so the entrypoint (the only place that owns the process) can arm the
+ * force-exit fallback.
+ */
+export interface MainCommandOutcome {
+  readonly headlessCompleted: boolean;
+}
+
+export async function handleMainCommand(
+  opts: CLIOptions,
+  version: string,
+): Promise<MainCommandOutcome> {
   let validated: ReturnType<typeof validateOptions>;
   try {
     validated = validateOptions(opts);
@@ -60,10 +76,11 @@ export async function handleMainCommand(opts: CLIOptions, version: string): Prom
 
   if (validated.uiMode === 'print') {
     await runPrompt(validated.options, version);
-    return;
+    return { headlessCompleted: true };
   }
 
   await runShell(validated.options, version);
+  return { headlessCompleted: false };
 }
 
 /** `kimi migrate`: launch the migration screen only, then exit. */
@@ -139,17 +156,32 @@ export function main(): void {
   const program = createProgram(
     version,
     (opts) => {
-      void handleMainCommand(opts, version).catch(async (error: unknown) => {
-        const operation = opts.prompt !== undefined ? 'run prompt' : 'start shell';
-        await logStartupFailure(operation, error);
-        process.stderr.write(
-          formatStartupError(error, {
-            operation,
-          }),
-        );
-        process.stderr.write(`See log: ${resolveGlobalLogPath(resolveKimiHome())}\n`);
-        process.exit(1);
-      });
+      void handleMainCommand(opts, version)
+        .then(async (outcome) => {
+          // Only the process entrypoint disposes of the process. Print mode
+          // relies on the event loop draining to exit; flush any buffered output
+          // and then arm an unref'd fallback so a stray ref'd handle left over
+          // from the run can't wedge a completed `kimi -p` until an external
+          // timeout. A healthy run drains and exits before the fallback fires.
+          if (outcome.headlessCompleted) {
+            await finalizeHeadlessRun(
+              process,
+              [process.stdout, process.stderr],
+              () => Number(process.exitCode) || 0,
+            );
+          }
+        })
+        .catch(async (error: unknown) => {
+          const operation = opts.prompt !== undefined ? 'run prompt' : 'start shell';
+          await logStartupFailure(operation, error);
+          process.stderr.write(
+            formatStartupError(error, {
+              operation,
+            }),
+          );
+          process.stderr.write(`See log: ${resolveGlobalLogPath(resolveKimiHome())}\n`);
+          process.exit(1);
+        });
     },
     () => {
       void handleMigrateCommand(version).catch(async (error: unknown) => {