diff --git a/.changeset/stale-hounds-resume.md b/.changeset/stale-hounds-resume.md new file mode 100644 index 0000000000..d66a09463f --- /dev/null +++ b/.changeset/stale-hounds-resume.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix a race where resuming a background subagent right after it was manually stopped could fail with an "already running" error. diff --git a/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts b/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts index c77d0485c0..edb98f91d7 100644 --- a/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts +++ b/packages/agent-core-v2/src/session/subagent/runAgentTurn.ts @@ -79,8 +79,10 @@ async function awaitRun( const controller = new AbortController(); const unlink = linkAbortSignal(options.signal, controller); const loop = target.accessor.get(IAgentLoopService); - const cancelTurn = (reason: unknown): void => { - loop.cancel(undefined, reason); + // Cancel by turn id: `loop.cancel(undefined, …)` only reaches the active + // turn, but this run's turn may still be queued when the abort lands. + const cancelTurn = (turnToCancel: Turn, reason: unknown): void => { + loop.cancel(turnToCancel.id, reason); }; let turnRef: Turn = turn; try { @@ -100,24 +102,43 @@ async function awaitRun( } finally { unlink(); if (controller.signal.aborted) { - cancelTurn(controller.signal.reason); + cancelTurn(turnRef, controller.signal.reason); } } } +/** + * Await a turn's terminal result, cancelling it first when the run aborts. + * + * This deliberately does NOT race `turn.result` against the abort signal: + * the loop only goes idle once `runTurn` unwinds (`releaseActiveTurn`), and + * downstream consumers — task settlement, the `task.killed` notification, + * and the resume guard reading `loop.status()` — must not observe the run + * as finished before that. A turn that never responds to the cancel is + * still bounded by the task layer's SIGTERM grace, not by an early + * rejection here. + */ async function awaitTurn( turn: Turn, controller: AbortController, - cancelTurn: (reason: unknown) => void, + cancelTurn: (turn: Turn, reason: unknown) => void, ): Promise { - const onAbort = (): void => { - cancelTurn(controller.signal.reason); + const cancelOnAbort = (): void => { + cancelTurn(turn, controller.signal.reason); }; - controller.signal.addEventListener('abort', onAbort, { once: true }); + controller.signal.addEventListener('abort', cancelOnAbort, { once: true }); try { - return await Promise.race([turn.result, abortPromise(controller.signal)]); + if (controller.signal.aborted) { + cancelOnAbort(); + } + const result = await turn.result; + // Rethrow the original abort reason instead of returning the cancelled + // result: consumers match the reason by identity (`isAbortError` / + // `error === sink.signal.reason`) to tell a kill apart from a failure. + controller.signal.throwIfAborted(); + return result; } finally { - controller.signal.removeEventListener('abort', onAbort); + controller.signal.removeEventListener('abort', cancelOnAbort); } } @@ -126,7 +147,7 @@ async function distillSummary( controller: AbortController, policy: AgentProfileSummaryPolicy | undefined, setTurn: (turn: Turn) => void, - cancelTurn: (reason: unknown) => void, + cancelTurn: (turn: Turn, reason: unknown) => void, ): Promise { const memory = target.accessor.get(IAgentContextMemoryService); let summary = latestAssistantText(memory.get()); @@ -194,21 +215,6 @@ function providerRateLimitErrorFromPayload(error: KimiErrorPayload): APIProvider return new APIProviderRateLimitError(error.message, requestId); } -function abortPromise(signal: AbortSignal): Promise { - if (signal.aborted) { - return Promise.reject(signal.reason ?? userCancellationReason()); - } - return new Promise((_resolve, reject) => { - signal.addEventListener( - 'abort', - () => { - reject(signal.reason ?? userCancellationReason()); - }, - { once: true }, - ); - }); -} - function latestAssistantText(messages: readonly ContextMessage[]): string { for (let i = messages.length - 1; i >= 0; i--) { const message = messages[i]!; diff --git a/packages/agent-core-v2/test/agent/task/idle-notification-repro.test.ts b/packages/agent-core-v2/test/agent/task/idle-notification-repro.test.ts index 67f4c547a0..115b5f0ac9 100644 --- a/packages/agent-core-v2/test/agent/task/idle-notification-repro.test.ts +++ b/packages/agent-core-v2/test/agent/task/idle-notification-repro.test.ts @@ -24,8 +24,11 @@ import { join } from 'pathe'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LifecycleScope, type IAgentScopeHandle } from '#/_base/di/scope'; +import type { generate as kosongGenerate } from '#/app/llmProtocol/generate'; import { IAgentTaskService } from '#/agent/task/task'; import { SubagentTask } from '#/session/subagent/tools/subagent-task'; +import { runAgentTurn } from '#/session/subagent/runAgentTurn'; import { IAgentProfileService } from '#/agent/profile/profile'; import { IAgentLoopService } from '#/agent/loop/loop'; import { @@ -238,6 +241,125 @@ describe('task notification → main agent (real Agent instance)', () => { }); }); + describe('kill ordering vs child loop unwind', () => { + type GenerateFn = typeof kosongGenerate; + + function agentScopeHandle(ctx: TestAgentContext, id: string): IAgentScopeHandle { + return { + id, + kind: LifecycleScope.Agent, + accessor: { get: ctx.get.bind(ctx) }, + dispose: () => {}, + } as IAgentScopeHandle; + } + + // Regression for: "manual stop of a background subagent → main + // auto-resumes → resume fails with 'already running and cannot run + // concurrently'". The killed task used to settle (and notify) the + // moment the abort landed — while the child loop was still unwinding, + // so the resume guard (`ensureOwnedIdleSubagent`, which reads + // `loop.status().state`) rejected the auto-resume. Settlement must + // wait for the loop to go idle. A turn that ignores the cancel stays + // bounded by the task layer's SIGTERM grace instead. + it('stop settles killed + notifies only after the child loop goes idle', async () => { + // Child agent whose in-flight LLM call unwinds slowly after cancel + // (models a tool mid-execution / a slow request abort): it rejects + // 200ms after the abort lands, not immediately. + let generateStarted!: () => void; + const inFlight = new Promise((resolve) => { + generateStarted = resolve; + }); + const slowToCancelGenerate: GenerateFn = async ( + _chat, + _systemPrompt, + _tools, + _history, + _callbacks, + options, + ) => { + const signal = options?.signal; + signal?.throwIfAborted(); + generateStarted(); + await new Promise((_resolve, reject) => { + signal?.addEventListener( + 'abort', + () => { + setTimeout(() => { + reject(signal.reason); + }, 200); + }, + { once: true }, + ); + }); + throw new Error('slowToCancelGenerate returned without being aborted'); + }; + + const main = createTestAgent(taskServices()); + const child = createTestAgent({ generate: slowToCancelGenerate }); + try { + const childHandle = agentScopeHandle(child, 'agent-child'); + const childLoop = child.get(IAgentLoopService); + + // Launch the subagent run (what AgentTool.launch does). + const controller = new AbortController(); + const run = await runAgentTurn( + childHandle, + { kind: 'prompt', prompt: 'do background work' }, + { signal: controller.signal }, + ); + // Mirror AgentTool.launch: the task handle maps summary → result. + const completion = run.completion.then((r) => ({ result: r.summary, usage: r.usage })); + void completion.catch(() => {}); + + // Wait until the in-flight step is genuinely parked inside the LLM + // call — the loop reports 'running' before the request starts, and + // stopping that early takes a different (already-fast) path. + await inFlight; + expect(childLoop.status().state).toBe('running'); + + const background = main.get(IAgentTaskService); + const taskId = background.registerTask( + new SubagentTask( + { agentId: 'agent-child', profileName: 'coder', completion }, + 'kill-order repro', + controller, + ), + { detached: true, timeoutMs: 0 }, + ); + + // The main agent is idle; the killed notification auto-launches a turn. + main.mockNextResponse({ type: 'text', text: 'ack from main agent' }); + const notificationTurnEnd = main.untilTurnEnd(); + + // Manual stop (TUI / REST path — no notification suppression). + const info = await background.stop(taskId, 'User initiated stop'); + expect(info?.status).toBe('killed'); + // Settlement waited for the child loop to unwind — this is the + // assertion the old race-based implementation fails. + expect(childLoop.status().state).toBe('idle'); + + // The task.killed notification reaches the main agent (this is what + // makes main call Agent(resume="agent-child")), and by then the + // resume guard's precondition already holds. + await vi.waitFor( + () => { + expect(main.llmCalls.length).toBeGreaterThanOrEqual(1); + }, + { timeout: 2000 }, + ); + const notified = JSON.stringify(main.llmCalls.at(-1)!.history); + expect(notified).toContain('task.killed'); + expect(notified).toContain(taskId); + expect(childLoop.status().state).toBe('idle'); + + await notificationTurnEnd; + } finally { + await main.dispose(); + await child.dispose(); + } + }); + }); + describe('resumed notifications', () => { let sessionDir: string; let ctx: TestAgentContext;