diff --git a/.changeset/print-drain-subagents.md b/.changeset/print-drain-subagents.md new file mode 100644 index 000000000..6d1e56480 --- /dev/null +++ b/.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. diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index 2d2c71b2d..66315ff3e 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/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 { diff --git a/apps/kimi-code/test/cli/run-prompt.test.ts b/apps/kimi-code/test/cli/run-prompt.test.ts index 924fcdad8..d75a6b3a5 100644 --- a/apps/kimi-code/test/cli/run-prompt.test.ts +++ b/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, }); }); diff --git a/packages/agent-core/src/agent/background/index.ts b/packages/agent-core/src/agent/background/index.ts index 5ef3fac69..9234668c8 100644 --- a/packages/agent-core/src/agent/background/index.ts +++ b/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 { + 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 { + 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 }, + ); + }); +} diff --git a/packages/agent-core/src/agent/index.ts b/packages/agent-core/src/agent/index.ts index fb759289d..131873155 100644 --- a/packages/agent-core/src/agent/index.ts +++ b/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; diff --git a/packages/agent-core/src/agent/turn/index.ts b/packages/agent-core/src/agent/turn/index.ts index b79b6f83b..db4904a34 100644 --- a/packages/agent-core/src/agent/turn/index.ts +++ b/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 ( diff --git a/packages/agent-core/src/rpc/core-api.ts b/packages/agent-core/src/rpc/core-api.ts index e397ae6ee..87f407b59 100644 --- a/packages/agent-core/src/rpc/core-api.ts +++ b/packages/agent-core/src/rpc/core-api.ts @@ -60,6 +60,7 @@ export interface CreateSessionPayload { readonly mcpServers?: Readonly>; readonly additionalDirs?: readonly string[]; readonly client?: ClientTelemetryInfo | undefined; + readonly drainAgentTasksOnStop?: boolean; } export interface CloseSessionPayload { diff --git a/packages/agent-core/src/rpc/core-impl.ts b/packages/agent-core/src/rpc/core-impl.ts index dfac26eb5..5866d0b46 100644 --- a/packages/agent-core/src/rpc/core-impl.ts +++ b/packages/agent-core/src/rpc/core-impl.ts @@ -296,6 +296,7 @@ export class KimiCore implements PromisableMethods { pluginCommands, appVersion: this.appVersion, additionalDirs, + drainAgentTasksOnStop: options.drainAgentTasksOnStop, }); try { session.metadata = { diff --git a/packages/agent-core/src/session/index.ts b/packages/agent-core/src/session/index.ts index da0b9b015..478c775b6 100644 --- a/packages/agent-core/src/session/index.ts +++ b/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; } diff --git a/packages/agent-core/test/agent/background/manager.test.ts b/packages/agent-core/test/agent/background/manager.test.ts index 7f74cefac..1b6afa031 100644 --- a/packages/agent-core/test/agent/background/manager.test.ts +++ b/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(): { + promise: Promise; + resolve: (value: T) => void; + reject: (reason?: unknown) => void; + } { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((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' }); + }); +}); diff --git a/packages/agent-core/test/agent/turn.test.ts b/packages/agent-core/test/agent/turn.test.ts index 417b2c5f7..852b95d42 100644 --- a/packages/agent-core/test/agent/turn.test.ts +++ b/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; @@ -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(() => {})) 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) }); diff --git a/packages/agent-core/test/session/lifecycle-hooks.test.ts b/packages/agent-core/test/session/lifecycle-hooks.test.ts index 38f55666c..e19badebd 100644 --- a/packages/agent-core/test/session/lifecycle-hooks.test.ts +++ b/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({ diff --git a/packages/node-sdk/src/types.ts b/packages/node-sdk/src/types.ts index 239189c37..82f1d1234 100644 --- a/packages/node-sdk/src/types.ts +++ b/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 {