feat: hold print-mode turn until background subagents drain (#1371)

* feat: hold print-mode turn until background subagents drain

In `kimi -p` (print mode), when the main agent ends a turn while background
subagents (`kind === 'agent'`) are still running, hold the turn open and
idle-wait until they finish, flushing their completions into the turn so the
model can react before the run exits.

Previously, the main agent could end its turn after launching background
subagents; the print flow then drained them with their completion
notifications suppressed, so the main agent never saw the results and the run
exited with the work abandoned (e.g. no nomination). This was the root cause
of the swarm-alpha-mining eval failures.

The hold is gated on a new `drainAgentTasksOnStop` session option (set by the
print flow), only affects `kind === 'agent'` background tasks, and is bounded
by `background.printWaitCeilingS`. Backfill / fan-out is handled by
re-enumerating active tasks. Other background task kinds and non-print modes
are unaffected.

.changeset/print-drain-subagents.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": minor
+---
+
+Wait for background subagents to finish and respond to their results before exiting in `kimi -p`, instead of ending the turn early.

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

@@ -338,6 +338,7 @@ async function resolvePromptSession(
     model,
     permission: 'auto',
     additionalDirs: opts.addDirs?.length ? opts.addDirs : undefined,
+    drainAgentTasksOnStop: true,
   });
   installHeadlessHandlers(session);
   return {

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

@@ -209,6 +209,7 @@ describe('runPrompt', () => {
       model: 'k2',
       permission: 'auto',
       additionalDirs: undefined,
+      drainAgentTasksOnStop: true,
     });
     expect(mocks.session.setPermission).not.toHaveBeenCalled();
     expect(mocks.session.setApprovalHandler).toHaveBeenCalledWith(expect.any(Function));
@@ -334,6 +335,7 @@ describe('runPrompt', () => {
       model: 'kimi-code/k2.5',
       permission: 'auto',
       additionalDirs: undefined,
+      drainAgentTasksOnStop: true,
     });
     expect(mocks.initializeTelemetry).toHaveBeenCalledWith(
       expect.objectContaining({ model: 'kimi-code/k2.5' }),
@@ -351,6 +353,7 @@ describe('runPrompt', () => {
       model: 'k2',
       permission: 'auto',
       additionalDirs: ['../shared', '/tmp/extra'],
+      drainAgentTasksOnStop: true,
     });
   });
 

packages/agent-core/src/agent/background/index.ts

@@ -534,6 +534,45 @@ export class BackgroundManager {
     return this.toInfo(entry);
   }
 
+  /**
+   * Wait until no active (non-terminal) task matching `predicate` remains.
+   *
+   * Used by print-mode (`kimi -p`) turn draining to hold a turn open while
+   * background subagents are still running. Re-enumerates after each batch
+   * settles so tasks registered during the wait (fan-out) are picked up.
+   * Resolves immediately when nothing matches. Bounded by `timeoutMs`; once
+   * the deadline passes the method returns without waiting for stragglers.
+   * Rejects with the abort reason when `signal` is aborted.
+   */
+  async waitForActiveTasks(
+    predicate: (info: BackgroundTaskInfo) => boolean,
+    options: { timeoutMs?: number; signal?: AbortSignal } = {},
+  ): Promise<void> {
+    const deadline =
+      options.timeoutMs !== undefined && options.timeoutMs > 0
+        ? Date.now() + options.timeoutMs
+        : undefined;
+    const signal = options.signal;
+    while (true) {
+      signal?.throwIfAborted();
+      const active = this.list(true).filter(predicate);
+      if (active.length === 0) return;
+      let perTaskTimeout: number | undefined;
+      if (deadline !== undefined) {
+        const remaining = deadline - Date.now();
+        if (remaining <= 0) return;
+        perTaskTimeout = remaining;
+      }
+      const batch = Promise.all(active.map((t) => this.wait(t.taskId, perTaskTimeout)));
+      if (signal === undefined) {
+        await batch;
+      } else {
+        await Promise.race([batch, abortRejecter(signal)]);
+      }
+      // Re-enumerate: settled tasks (and any fan-out) show up in the next list().
+    }
+  }
+
   /**
    * Wait until a foreground task either detaches from the current tool call or
    * reaches a terminal state. Detached tasks return immediately.
@@ -995,3 +1034,16 @@ function buildBackgroundTaskNotificationBody(info: BackgroundTaskInfo): string {
 
   return `${baseLine}${recovery}`;
 }
+
+function abortRejecter(signal: AbortSignal): Promise<never> {
+  if (signal.aborted) {
+    return Promise.reject(signal.reason ?? new Error('Aborted'));
+  }
+  return new Promise<never>((_, reject) => {
+    signal.addEventListener(
+      'abort',
+      () => reject(signal.reason ?? new Error('Aborted')),
+      { once: true },
+    );
+  });
+}

packages/agent-core/src/agent/index.ts

@@ -144,6 +144,22 @@ export class Agent {
   readonly goal: GoalMode;
   readonly replayBuilder: ReplayBuilder;
 
+  /**
+   * Print-mode (`kimi -p`) only: when true and the agent ends a turn while
+   * background subagents (`kind === 'agent'`) are still running, the turn loop
+   * holds the turn open and idle-waits until they finish, flushing their
+   * completions into the turn so the model can react before the run exits. Set
+   * by the session for print runs; defaults to false everywhere else.
+   */
+  printDrainAgentTasksOnStop = false;
+  /**
+   * Absolute deadline (ms epoch) bounding all print-mode drain waits for this
+   * agent, derived from `background.printWaitCeilingS`. `Infinity` when the
+   * drain is disabled. Shared across every drain hold within a single print
+   * run so backfill rounds cannot exceed the ceiling in aggregate.
+   */
+  printDrainDeadlineMs = Number.POSITIVE_INFINITY;
+
   private additionalDirs: readonly string[];
   private activeProfile?: ResolvedAgentProfile;
   private brandHome?: string;

packages/agent-core/src/agent/turn/index.ts

@@ -729,6 +729,28 @@ export class TurnFlow {
               if (this.flushSteerBuffer()) return { continue: true };
               signal.throwIfAborted();
 
+              // Print-mode drain: when `kimi -p` ends a turn while background
+              // subagents are still running, hold the turn open and idle-wait
+              // until they finish (or the drain deadline is reached). Their
+              // completions steer into the buffer during the wait and are
+              // flushed afterward, so the model gets one wrap-up step to react
+              // (nominate, backfill, ...) before the turn ends. Gated on a
+              // session flag so interactive / goal modes are unaffected.
+              if (this.agent.printDrainAgentTasksOnStop) {
+                const remaining = this.agent.printDrainDeadlineMs - Date.now();
+                const hasActiveAgentTask = this.agent.background
+                  .list(true)
+                  .some((task) => task.kind === 'agent');
+                if (hasActiveAgentTask && remaining > 0) {
+                  await this.agent.background.waitForActiveTasks(
+                    (task) => task.kind === 'agent',
+                    { timeoutMs: remaining, signal },
+                  );
+                  this.flushSteerBuffer();
+                  return { continue: true };
+                }
+              }
+
               // 2. After UpdateGoal marks a goal terminal, ask the model for one
               //    final user-facing outcome message before the turn ends.
               if (

packages/agent-core/src/rpc/core-api.ts

@@ -60,6 +60,7 @@ export interface CreateSessionPayload {
   readonly mcpServers?: Readonly<Record<string, McpServerConfig>>;
   readonly additionalDirs?: readonly string[];
   readonly client?: ClientTelemetryInfo | undefined;
+  readonly drainAgentTasksOnStop?: boolean;
 }
 
 export interface CloseSessionPayload {

packages/agent-core/src/rpc/core-impl.ts

@@ -296,6 +296,7 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
       pluginCommands,
       appVersion: this.appVersion,
       additionalDirs,
+      drainAgentTasksOnStop: options.drainAgentTasksOnStop,
     });
     try {
       session.metadata = {

packages/agent-core/src/session/index.ts

@@ -76,6 +76,12 @@ export interface SessionOptions {
   readonly appVersion?: string;
   readonly experimentalFlags?: ExperimentalFlagResolver;
   readonly additionalDirs?: readonly string[];
+  /**
+   * Print-mode (`kimi -p`) only: hold the main turn open while background
+   * subagents (`kind === 'agent'`) are still running, idle-waiting until they
+   * finish before the run exits. Set via the SDK `createSession` option.
+   */
+  readonly drainAgentTasksOnStop?: boolean;
 }
 
 export interface SessionSkillConfig {
@@ -293,6 +299,11 @@ export class Session {
     const { agent } = await this.createAgent({ type: 'main' }, {
       profile: DEFAULT_AGENT_PROFILES['agent'],
     });
+    if (this.options.drainAgentTasksOnStop) {
+      const ceilingS = this.options.background?.printWaitCeilingS ?? 3600;
+      agent.printDrainAgentTasksOnStop = true;
+      agent.printDrainDeadlineMs = Date.now() + ceilingS * 1000;
+    }
     await this.triggerSessionStart('startup');
     return agent;
   }

packages/agent-core/test/agent/background/manager.test.ts

@@ -15,6 +15,7 @@ import {
   BackgroundTaskPersistence,
   ProcessBackgroundTask,
   type BackgroundManager,
+  type BackgroundTaskInfo,
 } from '../../../src/agent/background';
 import {
   agentTask,
@@ -806,3 +807,96 @@ describe('BackgroundManager', () => {
     expect(await manager.readOutput(taskId)).toContain('bg-ok');
   }, 15_000);
 });
+
+
+describe('waitForActiveTasks', () => {
+  function deferred<T>(): {
+    promise: Promise<T>;
+    resolve: (value: T) => void;
+    reject: (reason?: unknown) => void;
+  } {
+    let resolve!: (value: T) => void;
+    let reject!: (reason?: unknown) => void;
+    const promise = new Promise<T>((res, rej) => {
+      resolve = res;
+      reject = rej;
+    });
+    return { promise, resolve, reject };
+  }
+
+  const isAgent = (info: BackgroundTaskInfo): boolean => info.kind === 'agent';
+
+  it('resolves immediately when no task matches the predicate', async () => {
+    const { manager } = createBackgroundManager();
+    // A process task does not match the agent predicate.
+    registerProcess(manager, immediateProcess(0), 'noop', 'proc');
+    await expect(manager.waitForActiveTasks(isAgent)).resolves.toBeUndefined();
+  });
+
+  it('waits until a matching agent task reaches a terminal state', async () => {
+    const { manager } = createBackgroundManager();
+    const done = deferred<{ result: string }>();
+    manager.registerTask(agentTask(done.promise, 'agent'));
+
+    let settled = false;
+    const wait = manager.waitForActiveTasks(isAgent).then(() => {
+      settled = true;
+    });
+    await new Promise((resolve) => setImmediate(resolve));
+    expect(settled).toBe(false);
+
+    done.resolve({ result: 'ok' });
+    await wait;
+    expect(settled).toBe(true);
+  });
+
+  it('re-enumerates tasks registered during the wait (fan-out)', async () => {
+    const { manager } = createBackgroundManager();
+    const first = deferred<{ result: string }>();
+    manager.registerTask(agentTask(first.promise, 'first'));
+
+    let settled = false;
+    const wait = manager.waitForActiveTasks(isAgent).then(() => {
+      settled = true;
+    });
+    await new Promise((resolve) => setImmediate(resolve));
+
+    // Fan out a second agent task after the wait started.
+    const second = deferred<{ result: string }>();
+    manager.registerTask(agentTask(second.promise, 'second'));
+
+    // Completing only the first must not settle the wait.
+    first.resolve({ result: '1' });
+    await new Promise((resolve) => setImmediate(resolve));
+    expect(settled).toBe(false);
+
+    second.resolve({ result: '2' });
+    await wait;
+    expect(settled).toBe(true);
+  });
+
+  it('returns when the timeout elapses even if a matching task is still running', async () => {
+    const { manager } = createBackgroundManager();
+    const done = deferred<{ result: string }>();
+    const taskId = manager.registerTask(agentTask(done.promise, 'stuck'));
+
+    await manager.waitForActiveTasks(isAgent, { timeoutMs: 20 });
+
+    // Task is still running (never resolved), but the wait returned on the deadline.
+    expect(manager.getTask(taskId)?.status).toBe('running');
+    done.resolve({ result: 'late' });
+  });
+
+  it('rejects when the signal is aborted', async () => {
+    const { manager } = createBackgroundManager();
+    const done = deferred<{ result: string }>();
+    manager.registerTask(agentTask(done.promise, 'agent'));
+    const controller = new AbortController();
+
+    const wait = manager.waitForActiveTasks(isAgent, { signal: controller.signal });
+    controller.abort(new Error('stop'));
+
+    await expect(wait).rejects.toThrow('stop');
+    done.resolve({ result: 'late' });
+  });
+});

packages/agent-core/test/agent/turn.test.ts

@@ -2,8 +2,10 @@ import { existsSync, mkdtempSync } from 'node:fs';
 import { tmpdir } from 'node:os';
 import { join } from 'pathe';
 import { setTimeout as delay } from 'node:timers/promises';
+import { Readable, type Writable } from 'node:stream';
 
-import type { Kaos } from '@moonshot-ai/kaos';
+import type { Kaos, KaosProcess } from '@moonshot-ai/kaos';
+import { createControlledPromise } from '@antfu/utils';
 import {
   APIConnectionError,
   APIEmptyResponseError,
@@ -18,6 +20,7 @@ import { describe, expect, it, vi } from 'vitest';
 import { HookEngine } from '../../src/session/hooks';
 import { abortError } from '../../src/utils/abort';
 import type { AgentOptions, AgentRecord, AgentRecordPersistence } from '../../src/agent';
+import { ProcessBackgroundTask } from '../../src/agent/background';
 import { InMemoryAgentRecordPersistence } from '../../src/agent/records';
 import { ErrorCodes, KimiError } from '../../src/errors';
 import type { Logger, LogPayload } from '../../src/logging';
@@ -30,6 +33,7 @@ import { recordingTelemetry, type TelemetryRecord } from '../fixtures/telemetry'
 import { createFakeKaos } from '../tools/fixtures/fake-kaos';
 import { createCommandKaos, testAgent, type TestAgentOptions } from './harness/agent';
 import { executeTool } from '../tools/fixtures/execute-tool';
+import { agentTask } from './background/helpers';
 
 type GenerateFn = NonNullable<AgentOptions['generate']>;
 
@@ -73,6 +77,66 @@ describe('Agent turn flow', () => {
     });
   });
 
+  it('holds the turn until a background subagent finishes, then runs a wrap-up step', async () => {
+    const ctx = testAgent();
+    ctx.agent.printDrainAgentTasksOnStop = true;
+    ctx.agent.printDrainDeadlineMs = Date.now() + 60_000;
+
+    const subDone = createControlledPromise<{ result: string }>();
+    ctx.agent.background.registerTask(agentTask(subDone, 'subagent'));
+
+    ctx.configure();
+    ctx.mockNextResponse({ type: 'text', text: 'first' });
+    ctx.mockNextResponse({ type: 'text', text: 'wrap-up' });
+
+    await ctx.rpc.prompt({ input: [{ type: 'text', text: 'go' }] });
+
+    let turnEnded = false;
+    const turnEnd = ctx.untilTurnEnd().then(() => {
+      turnEnded = true;
+    });
+
+    // Let the first model step finish and the drain hold engage.
+    for (let i = 0; i < 100 && ctx.llmCalls.length < 1; i++) await delay(5);
+    await delay(20);
+    expect(turnEnded).toBe(false);
+
+    // Completing the subagent releases the hold; the model takes a wrap-up step.
+    subDone.resolve({ result: 'sub-result' });
+    await turnEnd;
+
+    expect(turnEnded).toBe(true);
+    expect(ctx.llmCalls.length).toBe(2);
+  });
+
+  it('does not hold the turn for a non-agent (process) background task', async () => {
+    const ctx = testAgent();
+    ctx.agent.printDrainAgentTasksOnStop = true;
+    ctx.agent.printDrainDeadlineMs = Date.now() + 60_000;
+
+    const proc: KaosProcess = {
+      stdin: { write: vi.fn(), end: vi.fn() } as unknown as Writable,
+      stdout: Readable.from([]),
+      stderr: Readable.from([]),
+      pid: 4242,
+      exitCode: null,
+      wait: vi.fn().mockReturnValue(new Promise<number>(() => {})) as unknown as KaosProcess['wait'],
+      kill: vi.fn().mockResolvedValue(undefined) as unknown as KaosProcess['kill'],
+      dispose: vi.fn().mockResolvedValue(undefined) as unknown as KaosProcess['dispose'],
+    };
+    ctx.agent.background.registerTask(new ProcessBackgroundTask(proc, 'sleep 60', 'proc'));
+
+    ctx.configure();
+    ctx.mockNextResponse({ type: 'text', text: 'only step' });
+
+    await ctx.rpc.prompt({ input: [{ type: 'text', text: 'go' }] });
+    await ctx.untilTurnEnd();
+
+    // Process tasks do not trigger the subagent-only drain hold, so the turn
+    // ends after the single step.
+    expect(ctx.llmCalls.length).toBe(1);
+  });
+
   it('tracks turn_ended telemetry with protocol props', async () => {
     const records: TelemetryRecord[] = [];
     const ctx = testAgent({ telemetry: recordingTelemetry(records) });

packages/agent-core/test/session/lifecycle-hooks.test.ts

@@ -402,6 +402,42 @@ describe('Session lifecycle hooks', () => {
     expect(agent.background.getTask(taskId)?.status).toBe('killed');
   });
 
+  it('createMain enables print drain and sets a deadline when drainAgentTasksOnStop is true', async () => {
+    const { sessionDir, workDir } = await hookFixture();
+    const before = Date.now();
+    const session = new Session({
+      kaos: testKaos.withCwd(workDir),
+      id: 'session-print-drain',
+      homedir: sessionDir,
+      rpc: createSessionRpc(),
+      skills: { explicitDirs: [join(workDir, 'missing-skills')] },
+      background: { keepAliveOnExit: true, printWaitCeilingS: 42 },
+      drainAgentTasksOnStop: true,
+    });
+    const agent = await session.createMain();
+
+    expect(agent.printDrainAgentTasksOnStop).toBe(true);
+    expect(agent.printDrainDeadlineMs).toBeGreaterThanOrEqual(before + 42 * 1000 - 50);
+    expect(agent.printDrainDeadlineMs).toBeLessThanOrEqual(Date.now() + 42 * 1000 + 50);
+    await session.close();
+  });
+
+  it('createMain leaves print drain disabled by default', async () => {
+    const { sessionDir, workDir } = await hookFixture();
+    const session = new Session({
+      kaos: testKaos.withCwd(workDir),
+      id: 'session-print-drain-off',
+      homedir: sessionDir,
+      rpc: createSessionRpc(),
+      skills: { explicitDirs: [join(workDir, 'missing-skills')] },
+    });
+    const agent = await session.createMain();
+
+    expect(agent.printDrainAgentTasksOnStop).toBe(false);
+    expect(agent.printDrainDeadlineMs).toBe(Number.POSITIVE_INFINITY);
+    await session.close();
+  });
+
   it('cancels an active foreground turn before closing', async () => {
     const { sessionDir, workDir } = await hookFixture();
     const session = new Session({

packages/node-sdk/src/types.ts

@@ -104,6 +104,14 @@ export interface CreateSessionOptions {
   readonly persistenceKaos?: Kaos | undefined;
   readonly additionalDirs?: readonly string[];
   readonly sessionStartedProperties?: TelemetryProperties;
+  /**
+   * Print-mode (`kimi -p`) only: when the main agent ends a turn while
+   * background subagents (`kind === 'agent'`) are still running, hold the turn
+   * open and idle-wait until they all finish, flushing their completions into
+   * the turn so the model can react before the run exits. Ignored by
+   * interactive / SDK sessions.
+   */
+  readonly drainAgentTasksOnStop?: boolean;
 }
 
 export interface RenameSessionInput {