From 0235e0f297faa57192219aeaf7a4ee75f4d5ccc1 Mon Sep 17 00:00:00 2001 From: _Kerman Date: Thu, 9 Jul 2026 17:57:25 +0800 Subject: [PATCH] fix(kimi-code): drain v2 print subagents before exit --- .../kimi-code/src/cli/v2/create-v2-harness.ts | 7 +- apps/kimi-code/src/cli/v2/v2-session.ts | 68 +++++++- apps/kimi-code/test/cli/v2/v2-session.test.ts | 156 ++++++++++++++++-- 3 files changed, 215 insertions(+), 16 deletions(-) diff --git a/apps/kimi-code/src/cli/v2/create-v2-harness.ts b/apps/kimi-code/src/cli/v2/create-v2-harness.ts index cac3a6b23..87a9887c2 100644 --- a/apps/kimi-code/src/cli/v2/create-v2-harness.ts +++ b/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 { diff --git a/apps/kimi-code/src/cli/v2/v2-session.ts b/apps/kimi-code/src/cli/v2/v2-session.ts index f0629ea45..8fe5144d7 100644 --- a/apps/kimi-code/src/cli/v2/v2-session.ts +++ b/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 { @@ -115,6 +120,8 @@ export class V2Session implements PromptSession { } async waitForBackgroundTasksOnPrint(): Promise { + 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(TASK_CONFIG_SECTION) ?? - config.get(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(TASK_CONFIG_SECTION) ?? + config.get(LEGACY_BACKGROUND_CONFIG_SECTION) + ); + } + async createGoal(input: CreateGoalInput): Promise { return (await this.agent.accessor .get(IAgentGoalService) @@ -177,6 +209,34 @@ export class V2Session implements PromptSession { } } +async function waitForActiveAgentTasks( + taskService: IAgentTaskService, + signal: AbortSignal, +): Promise { + 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 { + if (signal.aborted) { + return Promise.reject(signal.reason ?? new Error('Aborted')); + } + return new Promise((_, 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 ''. diff --git a/apps/kimi-code/test/cli/v2/v2-session.test.ts b/apps/kimi-code/test/cli/v2/v2-session.test.ts index eca6a583d..516d2d66a 100644 --- a/apps/kimi-code/test/cli/v2/v2-session.test.ts +++ b/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 { @@ -63,15 +73,58 @@ function fakeAccessor(map: Map) { return { get: (token: unknown) => map.get(token) }; } -function buildSession(options: { ceilingS?: number; taskServices: FakeTaskService[] }): V2Session { +class FakeAfterStepSlot { + registration: + | { + id: string; + handler: (ctx: Record, next: () => Promise) => Promise; + options: unknown; + } + | undefined; + + register( + id: string, + handler: (ctx: Record, next: () => Promise) => Promise, + options?: unknown, + ) { + this.registration = { id, handler, options }; + return { dispose: () => {} }; + } + + async run(ctx: Record): Promise { + 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([ [ 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(); + 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([ [ 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); + }); +});