fix(kimi-code): drain v2 print subagents before exit

apps/kimi-code/src/cli/v2/create-v2-harness.ts

@@ -160,7 +160,12 @@ class V2PromptHarness implements PromptHarness {
       await agent.accessor.get(IAgentProfileService).setModel(options.model);
     }
     agent.accessor.get(IAgentPermissionModeService).setMode(options.permission ?? 'auto');
-    return new V2Session({ core: this.core, session, agent });
+    return new V2Session({
+      core: this.core,
+      session,
+      agent,
+      drainAgentTasksOnStop: options.drainAgentTasksOnStop,
+    });
   }
 
   async resumeSession(input: ResumeSessionInput): Promise<PromptSession> {

apps/kimi-code/src/cli/v2/v2-session.ts

@@ -9,8 +9,10 @@
 import {
   IAgentGoalService,
   IAgentLifecycleService,
+  IAgentLoopService,
   IAgentPermissionModeService,
   IAgentProfileService,
+  IAgentPromptService,
   IAgentPromptLegacyService,
   IAgentTaskService,
   IConfigService,
@@ -41,6 +43,7 @@ const TASK_CONFIG_SECTION = 'task';
 const LEGACY_BACKGROUND_CONFIG_SECTION = 'background';
 
 interface TaskPrintWaitConfig {
+  readonly keepAliveOnExit?: boolean;
   readonly printWaitCeilingS?: number;
 }
 
@@ -48,6 +51,7 @@ export interface V2SessionContext {
   readonly core: Scope;
   readonly session: ISessionScopeHandle;
   readonly agent: IAgentScopeHandle;
+  readonly drainAgentTasksOnStop?: boolean;
 }
 
 export class V2Session implements PromptSession {
@@ -64,6 +68,7 @@ export class V2Session implements PromptSession {
     this.agent = context.agent;
     this.id = context.session.id;
     this.workDir = resolveWorkDir(context.session);
+    if (context.drainAgentTasksOnStop === true) this.installPrintDrainHook();
   }
 
   async getStatus(): Promise<SessionStatus> {
@@ -115,6 +120,8 @@ export class V2Session implements PromptSession {
   }
 
   async waitForBackgroundTasksOnPrint(): Promise<void> {
+    if (!this.readKeepAliveOnExit()) return;
+
     // Drain background tasks (background bash and background subagents) spawned
     // during the turn before a `kimi -p` run exits. `-p` must be able to run
     // long tasks to completion, so we wait until every active task across every
@@ -154,11 +161,28 @@ export class V2Session implements PromptSession {
     if (allWaiters.length > 0) await Promise.all(allWaiters);
   }
 
+  private installPrintDrainHook(): void {
+    this.agent.accessor.get(IAgentPromptService);
+    const loop = this.agent.accessor.get(IAgentLoopService);
+    const tasks = this.agent.accessor.get(IAgentTaskService);
+    loop.hooks.afterStep.register(
+      'print-drain-agent-tasks',
+      async (ctx, next) => {
+        await next();
+        if (ctx.continue || ctx.finishReason === 'tool_calls') return;
+        if (await waitForActiveAgentTasks(tasks, ctx.signal)) ctx.continue = true;
+      },
+      { after: 'prompt-service-steer' },
+    );
+  }
+
+  private readKeepAliveOnExit(): boolean {
+    const section = this.readTaskPrintWaitConfig();
+    return section?.keepAliveOnExit === true;
+  }
+
   private readPrintWaitCeilingMs(): number {
-    const config = this.core.accessor.get(IConfigService);
-    const section =
-      config.get<TaskPrintWaitConfig>(TASK_CONFIG_SECTION) ??
-      config.get<TaskPrintWaitConfig>(LEGACY_BACKGROUND_CONFIG_SECTION);
+    const section = this.readTaskPrintWaitConfig();
     const ceilingS = section?.printWaitCeilingS;
     if (typeof ceilingS === 'number' && Number.isFinite(ceilingS) && ceilingS > 0) {
       return ceilingS * 1000;
@@ -166,6 +190,14 @@ export class V2Session implements PromptSession {
     return DEFAULT_PRINT_WAIT_CEILING_S * 1000;
   }
 
+  private readTaskPrintWaitConfig(): TaskPrintWaitConfig | undefined {
+    const config = this.core.accessor.get(IConfigService);
+    return (
+      config.get<TaskPrintWaitConfig>(TASK_CONFIG_SECTION) ??
+      config.get<TaskPrintWaitConfig>(LEGACY_BACKGROUND_CONFIG_SECTION)
+    );
+  }
+
   async createGoal(input: CreateGoalInput): Promise<GoalSnapshot> {
     return (await this.agent.accessor
       .get(IAgentGoalService)
@@ -177,6 +209,34 @@ export class V2Session implements PromptSession {
   }
 }
 
+async function waitForActiveAgentTasks(
+  taskService: IAgentTaskService,
+  signal: AbortSignal,
+): Promise<boolean> {
+  let waited = false;
+  while (true) {
+    signal.throwIfAborted();
+    const active = taskService.list(true).filter((task) => task.kind === 'agent');
+    if (active.length === 0) return waited;
+    waited = true;
+    const batch = Promise.all(active.map((task) => taskService.wait(task.taskId)));
+    await Promise.race([batch, abortRejecter(signal)]);
+  }
+}
+
+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 },
+    );
+  });
+}
+
 function resolveWorkDir(session: ISessionScopeHandle): string {
   // The session scope does not eagerly expose workDir; the index summary does
   // (cwd). We read it lazily from the session context if available, else ''.

apps/kimi-code/test/cli/v2/v2-session.test.ts

@@ -1,5 +1,7 @@
 import {
   IAgentLifecycleService,
+  IAgentLoopService,
+  IAgentPromptService,
   IAgentTaskService,
   IConfigService,
   type AgentTaskInfo,
@@ -13,6 +15,7 @@ import { V2Session } from '../../../src/cli/v2/v2-session';
 
 interface FakeTask {
   readonly taskId: string;
+  readonly kind?: AgentTaskInfo['kind'];
   /** ms until this task completes once `wait` is called on it. */
   readonly completesInMs: number;
   /** Optional task to spawn (append to the active list) when this task completes. */
@@ -29,7 +32,14 @@ class FakeTaskService {
   list(activeOnly?: boolean): readonly AgentTaskInfo[] {
     return this.tasks
       .filter((task) => !activeOnly || task.active)
-      .map((task) => ({ taskId: task.taskId, status: 'running' }) as unknown as AgentTaskInfo);
+      .map(
+        (task) =>
+          ({
+            taskId: task.taskId,
+            kind: task.kind ?? 'process',
+            status: 'running',
+          }) as unknown as AgentTaskInfo,
+      );
   }
 
   suppressTerminalNotification(taskId: string): Promise<void> {
@@ -63,15 +73,58 @@ function fakeAccessor(map: Map<unknown, unknown>) {
   return { get: (token: unknown) => map.get(token) };
 }
 
-function buildSession(options: { ceilingS?: number; taskServices: FakeTaskService[] }): V2Session {
+class FakeAfterStepSlot {
+  registration:
+    | {
+        id: string;
+        handler: (ctx: Record<string, unknown>, next: () => Promise<void>) => Promise<void>;
+        options: unknown;
+      }
+    | undefined;
+
+  register(
+    id: string,
+    handler: (ctx: Record<string, unknown>, next: () => Promise<void>) => Promise<void>,
+    options?: unknown,
+  ) {
+    this.registration = { id, handler, options };
+    return { dispose: () => {} };
+  }
+
+  async run(ctx: Record<string, unknown>): Promise<void> {
+    if (this.registration === undefined) return;
+    await this.registration.handler(ctx, async () => {});
+  }
+}
+
+class FakeLoopService {
+  readonly afterStep = new FakeAfterStepSlot();
+  readonly hooks = {
+    beforeStep: { register: () => ({ dispose: () => {} }) },
+    afterStep: this.afterStep,
+    onError: { register: () => ({ dispose: () => {} }) },
+  };
+}
+
+function buildSession(options: {
+  ceilingS?: number;
+  keepAliveOnExit?: boolean;
+  taskServices: FakeTaskService[];
+  drainAgentTasksOnStop?: boolean;
+  loop?: FakeLoopService;
+}): V2Session {
+  const taskConfig =
+    options.ceilingS !== undefined || options.keepAliveOnExit !== undefined
+      ? {
+          keepAliveOnExit: options.keepAliveOnExit,
+          printWaitCeilingS: options.ceilingS,
+        }
+      : undefined;
   const coreMap = new Map<unknown, unknown>([
     [
       IConfigService,
       {
-        get: (section: string) =>
-          section === 'task' && options.ceilingS !== undefined
-            ? { printWaitCeilingS: options.ceilingS }
-            : undefined,
+        get: (section: string) => (section === 'task' ? taskConfig : undefined),
       },
     ],
   ]);
@@ -81,6 +134,15 @@ function buildSession(options: { ceilingS?: number; taskServices: FakeTaskServic
     return { accessor: fakeAccessor(agentMap) } as unknown as IAgentScopeHandle;
   });
 
+  const mainAgentMap = new Map<unknown, unknown>();
+  if (options.taskServices[0] !== undefined) {
+    mainAgentMap.set(IAgentTaskService, options.taskServices[0]);
+  }
+  if (options.loop !== undefined) {
+    mainAgentMap.set(IAgentLoopService, options.loop);
+    mainAgentMap.set(IAgentPromptService, {});
+  }
+
   const sessionMap = new Map<unknown, unknown>([
     [
       IAgentLifecycleService,
@@ -93,26 +155,37 @@ function buildSession(options: { ceilingS?: number; taskServices: FakeTaskServic
   return new V2Session({
     core: { accessor: fakeAccessor(coreMap) } as unknown as Scope,
     session: { id: 'sess-1', accessor: fakeAccessor(sessionMap) } as unknown as ISessionScopeHandle,
-    agent: { id: 'main', accessor: fakeAccessor(new Map()) } as unknown as IAgentScopeHandle,
+    agent: { id: 'main', accessor: fakeAccessor(mainAgentMap) } as unknown as IAgentScopeHandle,
+    drainAgentTasksOnStop: options.drainAgentTasksOnStop,
   });
 }
 
 describe('V2Session.waitForBackgroundTasksOnPrint', () => {
   it('returns immediately when there are no active background tasks', async () => {
     const service = new FakeTaskService([]);
-    const session = buildSession({ taskServices: [service] });
+    const session = buildSession({ keepAliveOnExit: true, taskServices: [service] });
 
     await session.waitForBackgroundTasksOnPrint();
 
     expect(service.waitCalls).toHaveLength(0);
   });
 
-  it('waits for a background task to complete and bounds the wait by the default ceiling, not 30s', async () => {
+  it('returns immediately when keepAliveOnExit is not enabled', async () => {
     const service = new FakeTaskService([{ taskId: 'a', completesInMs: 20, active: true }]);
     const session = buildSession({ taskServices: [service] });
 
     await session.waitForBackgroundTasksOnPrint();
 
+    expect(service.waitCalls).toHaveLength(0);
+    expect(service.suppressed).toHaveLength(0);
+  });
+
+  it('waits for a background task to complete and bounds the wait by the default ceiling, not 30s', async () => {
+    const service = new FakeTaskService([{ taskId: 'a', completesInMs: 20, active: true }]);
+    const session = buildSession({ keepAliveOnExit: true, taskServices: [service] });
+
+    await session.waitForBackgroundTasksOnPrint();
+
     expect(service.waitCalls).toHaveLength(1);
     expect(service.waitCalls[0]?.taskId).toBe('a');
     // The old implementation hardcoded a 30s cap; the drain must use the 1h
@@ -125,7 +198,7 @@ describe('V2Session.waitForBackgroundTasksOnPrint', () => {
     const service = new FakeTaskService([
       { taskId: 'stuck', completesInMs: Number.POSITIVE_INFINITY, active: true },
     ]);
-    const session = buildSession({ ceilingS: 1, taskServices: [service] });
+    const session = buildSession({ ceilingS: 1, keepAliveOnExit: true, taskServices: [service] });
 
     const startedAt = Date.now();
     await session.waitForBackgroundTasksOnPrint();
@@ -142,7 +215,7 @@ describe('V2Session.waitForBackgroundTasksOnPrint', () => {
     const service = new FakeTaskService([
       { taskId: 'a', completesInMs: 20, active: true, spawnsOnComplete: spawned },
     ]);
-    const session = buildSession({ taskServices: [service] });
+    const session = buildSession({ keepAliveOnExit: true, taskServices: [service] });
 
     await session.waitForBackgroundTasksOnPrint();
 
@@ -152,3 +225,64 @@ describe('V2Session.waitForBackgroundTasksOnPrint', () => {
     expect(service.suppressed).toEqual(expect.arrayContaining(['a', 'b']));
   });
 });
+
+describe('V2Session print drain hook', () => {
+  it('waits for active background subagents before the print turn ends', async () => {
+    const loop = new FakeLoopService();
+    const spawned: FakeTask = {
+      taskId: 'agent-b',
+      kind: 'agent',
+      completesInMs: 20,
+      active: true,
+    };
+    const service = new FakeTaskService([
+      {
+        taskId: 'agent-a',
+        kind: 'agent',
+        completesInMs: 20,
+        active: true,
+        spawnsOnComplete: spawned,
+      },
+    ]);
+    buildSession({
+      taskServices: [service],
+      drainAgentTasksOnStop: true,
+      loop,
+    });
+
+    expect(loop.afterStep.registration?.id).toBe('print-drain-agent-tasks');
+    expect(loop.afterStep.registration?.options).toEqual({ after: 'prompt-service-steer' });
+    const ctx = {
+      signal: new AbortController().signal,
+      finishReason: 'completed',
+      continue: false,
+    };
+
+    await loop.afterStep.run(ctx);
+
+    expect(service.waitCalls.map((call) => call.taskId)).toEqual(['agent-a', 'agent-b']);
+    expect(ctx.continue).toBe(true);
+  });
+
+  it('does not hold the print turn for non-agent background tasks', async () => {
+    const loop = new FakeLoopService();
+    const service = new FakeTaskService([
+      { taskId: 'proc-a', kind: 'process', completesInMs: 20, active: true },
+    ]);
+    buildSession({
+      taskServices: [service],
+      drainAgentTasksOnStop: true,
+      loop,
+    });
+    const ctx = {
+      signal: new AbortController().signal,
+      finishReason: 'completed',
+      continue: false,
+    };
+
+    await loop.afterStep.run(ctx);
+
+    expect(service.waitCalls).toHaveLength(0);
+    expect(ctx.continue).toBe(false);
+  });
+});