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,25 @@ export async function runV2Print(
       setCrashPhase('shutdown');
       try {
         await restorePermission();
+        // A termination signal can arrive mid-turn: cancel turns and wait for idle agents first.
+        await raceWithTimeout(quiesceAgents(), CLI_SHUTDOWN_TIMEOUT_MS).catch(() => {});
       } finally {
-        if (telemetryService !== undefined) {
-          await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS);
+        try {
+          // Concurrent so the phases' allowances cannot sum past PROMPT_CLEANUP_TIMEOUT_MS.
+          await Promise.all([
+            // The turn's tail records reach the journal only via the wire's
+            // async persist queue; process.exit must not cut off that queue.
+            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 {
+          // Keep producers frozen until the journals are drained and disposed.
+          releaseQuiescence?.();
         }
-        await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }).catch(() => {});
-        app.dispose();
       }
     })());
     await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS);
@@ -254,6 +274,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 +893,117 @@ function countPendingBackgroundTasks(session: ISessionScopeHandle): number {
   return count;
 }
 
+/** Every agent handle in the session; the main agent is included explicitly since the lifecycle list skips `closing` agents. */
+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 producers, drain prompts, and cancel turns so closing records exist
+ * before the wire flush; returns a release holding a guard per loop.
+ */
+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.
+      return [];
+    }
+  });
+  // Task producers bypass the prompt queue and dispatch termination records
+  // straight to the wire; stop them first so the flush can persist those.
+  await Promise.allSettled(
+    handles.flatMap((handle) => {
+      try {
+        return [handle.accessor.get(IAgentTaskService).stopAllOnExit('Session closed')];
+      } catch {
+        return [];
+      }
+    }),
+  );
+  // Repeat until every queue is empty and every loop freezable: a prompt can
+  // still surface from the launch window or a cancelled turn's settle chain.
+  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. */
+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,