fix(kimi-code): flush wire journals before print-mode exit

.changeset/print-wire-flush-on-exit.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Fix print mode (`kimi -p`) losing session records when the run exits on an error or a termination signal.

apps/kimi-code/src/cli/v2/run-v2-print.ts

@@ -22,6 +22,7 @@ import {
   IAgentCronService,
   IAgentGoalService,
   IAgentLifecycleService,
+  IAgentLoopService,
   IAgentPermissionModeService,
   IAgentProfileService,
   IAgentPromptService,
@@ -30,6 +31,7 @@ import {
   IBootstrapService,
   IConfigService,
   IEventBus,
+  IEventDispatcher,
   IOAuthToolkit,
   ISessionIndex,
   ISessionManager,
@@ -122,6 +124,8 @@ import {
 const PROMPT_UI_MODE = 'print';
 /** Re-check `goalActive` at least this often while waiting for goal turns. */
 const GOAL_WAIT_POLL_MS = 250;
+/** Re-check each agent's prompt queue while waiting for it to drain at exit. */
+const PROMPT_QUIESCE_POLL_MS = 10;
 /**
  * Slack on top of a scheduled cron fire time while waiting for the steered
  * turn: covers the 1s tick poll interval plus fire → inject → turn-launch
@@ -196,6 +200,9 @@ export async function runV2Print(
   }
 
   let restorePermission = async (): Promise<void> => {};
+  let quiesceAgents = async (): Promise<void> => {};
+  let releaseQuiescence: (() => void) | undefined;
+  let flushWires = async (): Promise<void> => {};
   let removeTerminationCleanup: (() => void) | undefined;
   let cleanupPromise: Promise<void> | undefined;
   let telemetryService: ITelemetryService | undefined;
@@ -205,12 +212,39 @@ export async function runV2Print(
       setCrashPhase('shutdown');
       try {
         await restorePermission();
+        // A termination signal can arrive mid-turn: drain pending prompts,
+        // cancel any queued/active turns, and wait for every agent's loops
+        // and prompt queues to quiesce first, so the turn's cancellation and
+        // closing records exist before the flush below drains the journals.
+        // No-op when every agent is already idle.
+        await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {});
       } finally {
-        if (telemetryService !== undefined) {
-          await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS);
+        try {
+          // The shutdown phases are independent of each other; run them
+          // concurrently so their individual allowances cannot sum past the
+          // outer PROMPT_CLEANUP_TIMEOUT_MS bound and let the caller's
+          // process.exit cut off the tail (app.dispose included).
+          await Promise.all([
+            // The turn's tail records (step.end / turn.ended / prompt.completed)
+            // are dispatched fire-and-forget and reach the journal only through
+            // the wire service's async persist queue. Without an explicit flush,
+            // a cleanup that returns fast (e.g. telemetry disabled) lets
+            // process.exit cut off that queue before the records land on disk.
+            // Best-effort: a persist failure was already reported where the
+            // append failed.
+            raceWithTimeout(flushWires(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {}),
+            telemetryService !== undefined
+              ? raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS)
+              : Promise.resolve(),
+            shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }).catch(() => {}),
+          ]);
+          app.dispose();
+        } finally {
+          // Quiescence guards stay held across the flush and disposal so a
+          // late background-task completion or cron fire cannot start a new
+          // turn (and new records) after the journals were drained.
+          releaseQuiescence?.();
         }
-        await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }).catch(() => {});
-        app.dispose();
       }
     })());
     await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS);
@@ -254,6 +288,10 @@ export async function runV2Print(
 
     const resolved = await resolveNativeSession(app, opts, workDir, defaultModel, stderr);
     restorePermission = resolved.restorePermission;
+    quiesceAgents = async () => {
+      releaseQuiescence = await quiesceSessionAgents(resolved.session, resolved.agent);
+    };
+    flushWires = () => flushSessionWires(resolved.session, resolved.agent);
 
     telemetryService.setContext({ session_id: resolved.session.id, model: resolved.telemetryModel });
     setTelemetryContext({ sessionId: resolved.session.id });
@@ -869,6 +907,148 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number {
   return count;
 }
 
+/**
+ * Every agent handle in the session. The main agent handle is included
+ * explicitly: `IAgentLifecycleService` skips `closing` agents, but a closing
+ * agent's in-flight turn and already-dispatched tail records still deserve to
+ * land on disk.
+ */
+function collectSessionAgentHandles(
+  session: ISessionScopeHandle,
+  mainAgent: IAgentScopeHandle,
+): IAgentScopeHandle[] {
+  const agentManager = session.accessor.get(IAgentLifecycleService);
+  const handles = new Set<IAgentScopeHandle>([mainAgent]);
+  for (const agent of agentManager.list()) {
+    const handle = agentManager.handleOf(agent.agentId);
+    if (handle !== undefined) handles.add(handle);
+  }
+  return [...handles];
+}
+
+/**
+ * Stop every session agent's task producers, drain prompt queues, cancel
+ * queued and active turns, then wait for the loops to go idle and the prompt
+ * queues to empty. A termination signal can arrive mid-turn: without this,
+ * the turn's cancellation and closing records would only be produced by
+ * dispose()'s asynchronous teardown — after the wire flush, and after
+ * process.exit. Stopping and draining are no-ops on idle agents, so the
+ * normal (completed/failed turn) exit pays nothing here.
+ *
+ * On success returns a release function: every loop is left holding a
+ * quiescence guard so late producers (a cron fire enqueues straight into the
+ * loop, bypassing the prompt queue) cannot start a new turn — and new
+ * records — after the prompt queues read empty. The caller holds the release
+ * across the wire flush and disposal. Returns undefined when quiescence
+ * could not be reached within the caller's bound.
+ */
+async function quiesceSessionAgents(
+  session: ISessionScopeHandle,
+  mainAgent: IAgentScopeHandle,
+): Promise<(() => void) | undefined> {
+  const handles = collectSessionAgentHandles(session, mainAgent);
+  const promptServices = handles.flatMap((handle) => {
+    try {
+      return [handle.accessor.get(IAgentPromptService)];
+    } catch {
+      // A torn-down agent scope has no prompt service to drain or observe.
+      return [];
+    }
+  });
+  const loops = handles.flatMap((handle) => {
+    try {
+      return [handle.accessor.get(IAgentLoopService)];
+    } catch {
+      // A torn-down agent scope has no loop to quiesce; the wire flush below
+      // still covers its already-dispatched records.
+      return [];
+    }
+  });
+  // Task producers bypass the prompt queue and the loop guard below: their
+  // termination records are dispatched straight to the wire (and disposal
+  // would force-stop them after the flush). Stop them up front — settling a
+  // task dispatches its termination record now, so the flush can persist it
+  // (mirrors AgentLifecycleService.remove()).
+  await Promise.allSettled(
+    handles.flatMap((handle) => {
+      try {
+        return [handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed')];
+      } catch {
+        return [];
+      }
+    }),
+  );
+  // Repeat until a full pass finds every queue empty and every loop freezable:
+  // a prompt can surface after one pass already ran — from startNext's launch
+  // window (the prompt left `pending` and is not yet `active`, so drain could
+  // not see it) or from a cancelled turn's settle chain. The snapshot covers
+  // all three states; the loop reports idle (releaseActiveTurn) before the
+  // prompt-settle chain dispatches the final record, and settle() clears the
+  // active prompt and dispatches that record in one synchronous block, so an
+  // empty snapshot proves it was already queued for the wire flush.
+  for (;;) {
+    await Promise.allSettled(promptServices.map((service) => service.drain()));
+    for (const loop of loops) {
+      for (const turnId of loop.status().pendingTurnIds) loop.cancel(turnId);
+      loop.cancel();
+    }
+    await Promise.allSettled(loops.map((loop) => loop.settled()));
+    const guards: { dispose(): void }[] = [];
+    let frozen = true;
+    for (const loop of loops) {
+      let guard: { dispose(): void } | undefined;
+      try {
+        guard = loop.tryAcquireQuiescence();
+      } catch {
+        // A disposed loop cannot accept new submissions; it needs no guard.
+        continue;
+      }
+      if (guard === undefined) {
+        frozen = false;
+        break;
+      }
+      guards.push(guard);
+    }
+    const busy = promptServices.some((service) => {
+      try {
+        const snapshot = service.list();
+        return (
+          snapshot.launching ||
+          snapshot.active !== undefined ||
+          snapshot.pending.length > 0
+        );
+      } catch {
+        return false;
+      }
+    });
+    if (frozen && !busy) {
+      return () => {
+        for (const guard of guards) guard.dispose();
+      };
+    }
+    for (const guard of guards) guard.dispose();
+    await new Promise((resolve) => {
+      setTimeout(resolve, PROMPT_QUIESCE_POLL_MS);
+    });
+  }
+}
+
+/**
+ * Flush every session agent's wire journal. Each flush settles independently:
+ * one agent's broken journal must not cut the wait short for the others (the
+ * caller proceeds to `process.exit`).
+ */
+async function flushSessionWires(
+  session: ISessionScopeHandle,
+  mainAgent: IAgentScopeHandle,
+): Promise<void> {
+  await Promise.allSettled(
+    collectSessionAgentHandles(session, mainAgent).map((handle) =>
+      handle.accessor.get(IEventDispatcher).flush(),
+    ),
+  );
+}
+
 async function drainBackgroundTasks(
   session: ISessionScopeHandle,
   ceilingS: number | undefined,