feat: hold print-mode turn until background subagents drain

.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) });