From 7e47efe0d5f1a10a9810ec48a503dd646188fe83 Mon Sep 17 00:00:00 2001 From: qer Date: Tue, 14 Jul 2026 05:15:56 +0800 Subject: [PATCH 1/4] fix: stop individual swarm members --- .changeset/web-stop-swarm-members.md | 5 + apps/kimi-web/src/App.vue | 1 + .../src/api/daemon/agentEventProjector.ts | 7 +- .../components/chat/tool-calls/SwarmTool.vue | 59 ++- apps/kimi-web/src/lib/swarmCardRows.ts | 4 + .../test/agent-event-projector.test.ts | 53 ++ apps/kimi-web/test/swarm-card-rows.test.ts | 4 +- apps/kimi-web/test/workspace-state.test.ts | 9 +- packages/agent-core-v2/src/agent/task/task.ts | 5 +- .../src/agent/task/taskService.ts | 32 +- .../session/agentLifecycle/mirrorAgentRun.ts | 24 +- .../src/session/swarm/agentRunBatch.ts | 66 ++- .../src/session/swarm/sessionSwarm.ts | 23 + .../src/session/swarm/sessionSwarmService.ts | 294 +++++++++-- .../test/agent/swarm/swarm.test.ts | 9 +- .../test/agent/task/taskManager.test.ts | 48 ++ .../test/agent/task/tools/task-tools.test.ts | 6 + packages/agent-core-v2/test/harness/agent.ts | 1 + .../node-local/hostProcessService.test.ts | 123 ++++- .../os/backends/node-local/tools/bash.test.ts | 6 + .../test/session/swarm/sessionSwarm.test.ts | 486 +++++++++++++++++- packages/agent-core-v2/test/tool/tool.test.ts | 49 ++ packages/kap-server/src/routes/tasks.ts | 103 +++- packages/kap-server/test/tasks.test.ts | 152 +++++- .../protocol/src/__tests__/events.test.ts | 17 + packages/protocol/src/events.ts | 3 + 26 files changed, 1470 insertions(+), 119 deletions(-) create mode 100644 .changeset/web-stop-swarm-members.md diff --git a/.changeset/web-stop-swarm-members.md b/.changeset/web-stop-swarm-members.md new file mode 100644 index 0000000000..40adb1d9de --- /dev/null +++ b/.changeset/web-stop-swarm-members.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +web: Allow active swarm members to be stopped individually, including cancellation of their running shell commands. diff --git a/apps/kimi-web/src/App.vue b/apps/kimi-web/src/App.vue index b5167633fe..9aa32680a1 100644 --- a/apps/kimi-web/src/App.vue +++ b/apps/kimi-web/src/App.vue @@ -71,6 +71,7 @@ provide( 'resolveSwarmMembers', (toolCallId: string): SwarmMember[] => client.swarmMembersByToolCallId.value.get(toolCallId) ?? [], ); +provide('cancelSubagent', (agentId: string): Promise => client.cancelTask(agentId)); const { t } = useI18n(); // KAP/daemon debug panel — opt-in via ?debug=1 or localStorage kimi-web.debug=1. diff --git a/apps/kimi-web/src/api/daemon/agentEventProjector.ts b/apps/kimi-web/src/api/daemon/agentEventProjector.ts index bcfefbb5d3..660006b3e6 100644 --- a/apps/kimi-web/src/api/daemon/agentEventProjector.ts +++ b/apps/kimi-web/src/api/daemon/agentEventProjector.ts @@ -1065,9 +1065,10 @@ export function createAgentProjector(): AgentProjector { case 'subagent.failed': { const outputPreview = typeof p?.error === 'string' ? p.error : undefined; + const status = p?.cancelled === true ? 'cancelled' : 'failed'; const task = patchSubagent(s, sessionId, p?.subagentId, { subagentPhase: 'failed', - status: 'failed', + status, completedAt: new Date().toISOString(), outputPreview, }); @@ -1076,7 +1077,7 @@ export function createAgentProjector(): AgentProjector { type: 'taskCompleted', sessionId, taskId: p?.subagentId ?? '', - status: 'failed', + status, outputPreview, }); break; @@ -1152,7 +1153,7 @@ export function createAgentProjector(): AgentProjector { : typeof info.taskId === 'number' ? String(info.taskId) : '', - status: failed ? 'failed' : 'completed', + status: info.status === 'killed' ? 'cancelled' : failed ? 'failed' : 'completed', // Do NOT set outputPreview here. The command is already kept on the // task as `command`; setting outputPreview to `$ ` would // clobber any real output captured by polling and prevents the UI diff --git a/apps/kimi-web/src/components/chat/tool-calls/SwarmTool.vue b/apps/kimi-web/src/components/chat/tool-calls/SwarmTool.vue index 15e0ce2ad2..cf7e644829 100644 --- a/apps/kimi-web/src/components/chat/tool-calls/SwarmTool.vue +++ b/apps/kimi-web/src/components/chat/tool-calls/SwarmTool.vue @@ -17,6 +17,7 @@ import { toolLabel } from '../../../lib/toolMeta'; import { parseSwarmResult } from '../../../lib/parseSwarmResult'; import { buildSwarmCardRows, type SwarmCardRow } from '../../../lib/swarmCardRows'; import Icon from '../../ui/Icon.vue'; +import IconButton from '../../ui/IconButton.vue'; import StatusDot from '../../ui/StatusDot.vue'; import Tooltip from '../../ui/Tooltip.vue'; @@ -60,6 +61,7 @@ function parseInput(arg: string): SwarmInput { const resolveSwarmMembers = inject<(toolCallId: string) => SwarmMember[] | undefined>('resolveSwarmMembers'); +const cancelSubagent = inject<(agentId: string) => Promise>('cancelSubagent'); const input = computed(() => parseInput(props.tool.arg)); const label = computed(() => toolLabel(props.tool.name)); @@ -151,6 +153,10 @@ function isRowOpen(id: string): boolean { function phaseLabel(phase: AppSubagentPhase): string { return t(`tools.swarm.phase${phase[0]!.toUpperCase()}${phase.slice(1)}`); } + +function stopSubagent(agentId: string): void { + void cancelSubagent?.(agentId); +} @@ -395,11 +412,16 @@ function phaseLabel(phase: AppSubagentPhase): string { .member:last-child { border-bottom: none; } +.member-line { + display: flex; + align-items: center; +} .member-head { display: flex; align-items: center; gap: 8px; - width: 100%; + flex: 1; + min-width: 0; min-height: 32px; padding: 0 11px; border: none; @@ -456,6 +478,9 @@ function phaseLabel(phase: AppSubagentPhase): string { color: var(--color-text-faint); flex: none; } +.member-stop { + margin-right: var(--space-1); +} .member-body { padding: 4px 11px 10px 31px; color: var(--color-text-muted); diff --git a/apps/kimi-web/src/lib/swarmCardRows.ts b/apps/kimi-web/src/lib/swarmCardRows.ts index f4d9458704..d250086311 100644 --- a/apps/kimi-web/src/lib/swarmCardRows.ts +++ b/apps/kimi-web/src/lib/swarmCardRows.ts @@ -14,6 +14,8 @@ export interface SwarmCardRow { activity: string; phase: AppSubagentPhase; body: string; + /** Only live, non-terminal members have an Agent id accepted by stop. */ + canStop: boolean; } function lastNonEmptyLine(text: string | undefined): string { @@ -53,6 +55,7 @@ function resultRow(sub: SwarmResultSubagent, index: number): SwarmCardRow { activity: sub.body.split('\n')[0] ?? '', phase: outcomeToPhase(sub.outcome), body: sub.body, + canStop: false, }; } @@ -85,6 +88,7 @@ export function buildSwarmCardRows(members: SwarmMember[], result: SwarmResult | activity: swarmMemberActivity(m), phase: m.phase, body: swarmMemberBody(m), + canStop: m.phase === 'queued' || m.phase === 'working' || m.phase === 'suspended', })); if (!result) return memberRows; diff --git a/apps/kimi-web/test/agent-event-projector.test.ts b/apps/kimi-web/test/agent-event-projector.test.ts index 36bf5ec581..a66402eac6 100644 --- a/apps/kimi-web/test/agent-event-projector.test.ts +++ b/apps/kimi-web/test/agent-event-projector.test.ts @@ -66,6 +66,59 @@ describe('subagent streaming text', () => { }); }); +describe('subagent cancellation', () => { + it('projects an explicitly cancelled subagent as cancelled', () => { + const projector = createAgentProjector(); + projector.project( + 'subagent.spawned', + { agentId: 'main', subagentId: 'agent-child', runInBackground: false }, + 's1', + ); + + const events = projector.project( + 'subagent.failed', + { + agentId: 'main', + subagentId: 'agent-child', + error: 'Aborted by the user', + cancelled: true, + }, + 's1', + ); + + expect(events).toContainEqual( + expect.objectContaining({ + type: 'taskCompleted', + taskId: 'agent-child', + status: 'cancelled', + }), + ); + }); + + it('maps a killed task termination to cancelled', () => { + const projector = createAgentProjector(); + const events = projector.project( + 'task.terminated', + { + agentId: 'main', + info: { + taskId: 'task-1', + status: 'killed', + }, + }, + 's1', + ); + + expect(events).toEqual([ + expect.objectContaining({ + type: 'taskCompleted', + taskId: 'task-1', + status: 'cancelled', + }), + ]); + }); +}); + describe('agent error projection', () => { it('drops a subagent error instead of surfacing it as a session warning', () => { const projector = createAgentProjector(); diff --git a/apps/kimi-web/test/swarm-card-rows.test.ts b/apps/kimi-web/test/swarm-card-rows.test.ts index 957e8e56e0..16943e67bb 100644 --- a/apps/kimi-web/test/swarm-card-rows.test.ts +++ b/apps/kimi-web/test/swarm-card-rows.test.ts @@ -64,7 +64,7 @@ describe('buildSwarmCardRows', () => { [member('a', '子任务 A', { text: 'streaming' })], null, ); - expect(rows).toEqual([{ id: 'a', name: '子任务 A', activity: 'streaming', phase: 'working', body: 'streaming' }]); + expect(rows).toEqual([{ id: 'a', name: '子任务 A', activity: 'streaming', phase: 'working', body: 'streaming', canStop: true }]); }); it('builds rows from result subagents when no members are present', () => { @@ -77,6 +77,7 @@ describe('buildSwarmCardRows', () => { ); expect(rows.map((r) => r.name)).toEqual(['A', 'B']); expect(rows.map((r) => r.phase)).toEqual(['completed', 'failed']); + expect(rows.every((r) => r.canStop === false)).toBe(true); }); it('appends result-only aborted not_started rows on top of live members', () => { @@ -94,6 +95,7 @@ describe('buildSwarmCardRows', () => { expect(rows.map((r) => r.id)).toEqual(['a1', 'a2', 'C']); expect(rows[2]?.phase).toBe('failed'); expect(rows[2]?.body).toBe('C never started'); + expect(rows[2]?.canStop).toBe(false); }); it('does not duplicate a result row that a live member already covers', () => { diff --git a/apps/kimi-web/test/workspace-state.test.ts b/apps/kimi-web/test/workspace-state.test.ts index 6e5c4c6821..76a7021416 100644 --- a/apps/kimi-web/test/workspace-state.test.ts +++ b/apps/kimi-web/test/workspace-state.test.ts @@ -562,15 +562,18 @@ describe('useWorkspaceState — cancelTask', () => { expect(state.tasksBySession['sess_1']?.[0]?.status).toBe('running'); }); - it('marks the task cancelled on success', async () => { + it('sends the displayed Agent id and marks the subagent cancelled on success', async () => { apiMock.cancelTask.mockResolvedValue({ cancelled: true }); const state = createState(); - state.tasksBySession = { sess_1: [task('t_1', 'running')] }; + state.tasksBySession = { + sess_1: [{ ...task('agent_1', 'running'), kind: 'subagent' }], + }; const deps = createDeps(); const ws = useWorkspaceState(state, deps); - await ws.cancelTask('t_1'); + await ws.cancelTask('agent_1'); + expect(apiMock.cancelTask).toHaveBeenCalledWith('sess_1', 'agent_1'); expect(state.tasksBySession['sess_1']?.[0]?.status).toBe('cancelled'); expect(deps.pushOperationFailure).not.toHaveBeenCalled(); }); diff --git a/packages/agent-core-v2/src/agent/task/task.ts b/packages/agent-core-v2/src/agent/task/task.ts index a0c5feb824..c22e6058f5 100644 --- a/packages/agent-core-v2/src/agent/task/task.ts +++ b/packages/agent-core-v2/src/agent/task/task.ts @@ -5,7 +5,9 @@ * detached work. Task execution adapters implement the generic `AgentTask` * contract from this domain's type module; this service owns registration, * output retention, persistence, detach/stop/wait, and terminal notifications. - * Bound at Agent scope. + * Agent tasks can also be resolved by their stable Agent id so protocol edges + * do not have to confuse that identity with the generated task id. Bound at + * Agent scope. */ import { createDecorator } from '#/_base/di/instantiation'; @@ -108,6 +110,7 @@ export interface IAgentTaskService { /** @deprecated Use `taskService.run()` + `track()` instead. */ registerTask(task: AgentTask, options?: RegisterAgentTaskOptions): string; getTask(taskId: string): AgentTaskInfo | undefined; + getAgentTask(agentId: string): AgentTaskInfo | undefined; list(activeOnly?: boolean, limit?: number): readonly AgentTaskInfo[]; persistOutput(taskId: string): void; getOutputSnapshot( diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index e44e87cbdd..ba514d0925 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -24,7 +24,7 @@ import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; import type { ContentPart } from '#/app/llmProtocol/message'; import { Disposable } from '#/_base/di/lifecycle'; -import { abortable } from '#/_base/utils/abort'; +import { abortable, userCancellationReason } from '#/_base/utils/abort'; import { escapeXml, escapeXmlAttr } from '#/_base/utils/xml-escape'; import { IEventBus } from '#/app/event/eventBus'; import type { ContextMessage, TaskOrigin } from '#/agent/contextMemory/types'; @@ -455,6 +455,17 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { return entry === undefined ? this.ghosts.get(taskId) : this.toInfo(entry); } + getAgentTask(agentId: string): AgentTaskInfo | undefined { + let found: AgentTaskInfo | undefined; + for (const entry of this.tasks.values()) { + found = newerAgentTaskForAgent(found, this.toInfo(entry), agentId); + } + for (const ghost of this.ghosts.values()) { + found = newerAgentTaskForAgent(found, ghost, agentId); + } + return found; + } + list(activeOnly = true, limit?: number): readonly AgentTaskInfo[] { const result: AgentTaskInfo[] = []; for (const entry of this.tasks.values()) { @@ -648,7 +659,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { const normalized = normalizeReason(reason); return this.terminateWithGrace(entry, { stopReason: normalized, - abortReason: normalized, + abortReason: userCancellationReason(), finalStatus: 'killed', }); } @@ -1211,6 +1222,23 @@ function newerRestoredTask( return loaded; } +function newerAgentTaskForAgent( + current: AgentTaskInfo | undefined, + candidate: AgentTaskInfo, + agentId: string, +): AgentTaskInfo | undefined { + if (candidate.kind !== 'agent' || candidate.agentId !== agentId) return current; + if (current === undefined) return candidate; + + const currentActive = !isAgentTaskTerminal(current.status); + const candidateActive = !isAgentTaskTerminal(candidate.status); + if (currentActive !== candidateActive) return candidateActive ? candidate : current; + + const currentTimestamp = current.endedAt ?? current.startedAt; + const candidateTimestamp = candidate.endedAt ?? candidate.startedAt; + return candidateTimestamp >= currentTimestamp ? candidate : current; +} + type TaskNotificationOrigin = Pick; function isTaskOrigin(origin: unknown): origin is TaskNotificationOrigin { diff --git a/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts b/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts index 1043f59261..49f71d37f4 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts @@ -19,10 +19,16 @@ * completed / failed` and telemetry still tracks `subagent_created` so existing * session recordings and dashboards stay valid. Rename lives on a separate * wire-cleanup PR. + * Explicit cancellation remains observable as `subagent.failed` with the + * optional cancellation marker even when retry-related failures are hidden. */ import type { IAgentScopeHandle } from '#/_base/di/scope'; -import { userCancellationReason } from '#/_base/utils/abort'; +import { + isAbortError, + isUserCancellation, + userCancellationReason, +} from '#/_base/utils/abort'; import { IAgentContextSizeService } from '#/agent/contextSize/contextSize'; import { isProviderRateLimitError } from '#/app/llmProtocol/errors'; import { type TokenUsage } from '#/app/llmProtocol/usage'; @@ -34,7 +40,6 @@ import type { SubagentStartedEvent, } from '@moonshot-ai/protocol'; import { IEventBus } from '#/app/event/eventBus'; -import { isAbortError } from '#/_base/utils/abort'; import { type AgentRunHandle, IAgentLifecycleService } from './agentLifecycle'; @@ -65,7 +70,6 @@ export interface MirrorAgentRunOptions { * retry turns, which skip the hook. */ readonly prompt?: string; - /** Skip the requester-side `subagent.failed` record for provider-rate-limit / aborted failures. */ readonly suppressRateLimitFailureEvent?: boolean; /** The requester's cancellation signal (passed through to the start hook slot). */ readonly signal: AbortSignal; @@ -150,7 +154,19 @@ export async function mirrorAgentRun( }); return result; } catch (error) { - if (!isAbortError(error) && !shouldSuppressFailure(options, error)) { + const cancellationReason = isUserCancellation(error) + ? error + : isUserCancellation(options.signal.reason) + ? options.signal.reason + : undefined; + if (cancellationReason !== undefined) { + eventBus?.publish({ + type: 'subagent.failed', + subagentId: run.agentId, + error: errorMessage(cancellationReason), + cancelled: true, + }); + } else if (!isAbortError(error) && !shouldSuppressFailure(options, error)) { eventBus?.publish({ type: 'subagent.failed', subagentId: run.agentId, diff --git a/packages/agent-core-v2/src/session/swarm/agentRunBatch.ts b/packages/agent-core-v2/src/session/swarm/agentRunBatch.ts index de06f56646..7bfe73710a 100644 --- a/packages/agent-core-v2/src/session/swarm/agentRunBatch.ts +++ b/packages/agent-core-v2/src/session/swarm/agentRunBatch.ts @@ -5,15 +5,20 @@ * loop used by `SessionSwarmService`; drives each attempt through a * `AgentRunBatchLauncher` and surfaces requeues via `suspended`. Pure scheduling * logic — owns no scoped state. Not part of the public surface: only - * `SessionSwarmService` imports it. + * `SessionSwarmService` imports it. Supports cancellation of one identified + * member without stopping unrelated work in the batch. */ import { isProviderRateLimitError } from '#/app/llmProtocol/errors'; import { type TokenUsage } from '#/app/llmProtocol/usage'; import * as retry from 'retry'; -import { isUserCancellation } from '#/_base/utils/abort'; -import type { SessionSwarmRunResult, SessionSwarmTask } from './sessionSwarm'; +import { isUserCancellation, userCancellationReason } from '#/_base/utils/abort'; +import type { + SessionSwarmRunResult, + SessionSwarmStopResult, + SessionSwarmTask, +} from './sessionSwarm'; // ── Launcher contract ──────────────────────────────────────────────── // @@ -31,6 +36,7 @@ export interface AgentRunAttemptOptions { readonly swarmIndex?: number; readonly runInBackground: boolean; readonly signal: AbortSignal; + readonly onAgentIdentified?: (agentId: string) => void; readonly onReady?: () => void; readonly suppressRateLimitFailureEvent?: boolean; } @@ -210,6 +216,49 @@ export class AgentRunBatch { }); } + stopAgent(agentId: string): SessionSwarmStopResult { + const terminal = this.results.find((result) => result?.agentId === agentId); + if (terminal !== undefined) { + return { + kind: 'already_terminal', + agentId, + status: terminal.status, + }; + } + + const active = Array.from(this.active).find( + (attempt) => attempt.state.agentId === agentId, + ); + if (active !== undefined) { + if (!active.controller.signal.aborted) { + active.controller.abort(userCancellationReason()); + } + return { kind: 'stopping', agentId }; + } + + const suspendedIndex = this.pending.findIndex( + (state) => state.agentId === agentId && state.retryAgentId === agentId, + ); + if (suspendedIndex !== -1) { + const [state] = this.pending.splice(suspendedIndex, 1); + if (state === undefined) return { kind: 'not_found', agentId }; + this.results[state.index] = { + task: state.task, + agentId, + status: 'aborted', + state: 'started', + error: 'The user manually interrupted this subagent batch.', + }; + this.schedule(); + return { + kind: 'stopped', + agentId, + }; + } + + return { kind: 'not_found', agentId }; + } + private schedule(): void { if (this.finished) return; if (this.finishIfComplete()) return; @@ -316,6 +365,11 @@ export class AgentRunBatch { swarmIndex: task.swarmIndex, runInBackground: task.runInBackground, signal: attempt.controller.signal, + onAgentIdentified: (agentId) => { + if (!this.finished && this.active.has(attempt)) { + attempt.state.agentId = agentId; + } + }, onReady: () => { this.markAttemptReady(attempt); }, @@ -352,6 +406,9 @@ export class AgentRunBatch { usage: completion.usage, }; } catch (error) { + if (attempt.controller.signal.aborted) { + return this.failedAttemptOutcome(attempt, error); + } if (isProviderRateLimitError(error)) { return { type: 'rate_limited', @@ -397,7 +454,6 @@ export class AgentRunBatch { private handleAttemptOutcome(attempt: ActiveAttempt, outcome: AttemptOutcome): void { if (!this.releaseAttempt(attempt)) return; if (this.finished) return; - if ('status' in outcome) { this.results[attempt.state.index] = outcome; } else if (this.isOnlyUnfinishedTask(attempt.state)) { @@ -687,5 +743,3 @@ export function resolveSwarmMaxConcurrency( } return value; } - - diff --git a/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts b/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts index b092b2cfec..1ad64f0aa1 100644 --- a/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts +++ b/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts @@ -52,6 +52,25 @@ export interface SessionSwarmRunResult { readonly error?: string; } +export type SessionSwarmStopResult = + | { + readonly kind: 'stopping'; + readonly agentId: string; + } + | { + readonly kind: 'stopped'; + readonly agentId: string; + } + | { + readonly kind: 'already_terminal'; + readonly agentId: string; + readonly status: SessionSwarmRunResult['status']; + } + | { + readonly kind: 'not_found'; + readonly agentId: string; + }; + export interface ISessionSwarmService { readonly _serviceBrand: undefined; @@ -60,6 +79,10 @@ export interface ISessionSwarmService { readonly agentId: string; }): Promise; run(args: SessionSwarmRunArgs): Promise[]>; + stopAgent(args: { + readonly callerAgentId: string; + readonly agentId: string; + }): SessionSwarmStopResult; cancel(args: { readonly callerAgentId: string }): void; } diff --git a/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts b/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts index e7d3e30534..d1d5f5987a 100644 --- a/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts +++ b/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts @@ -4,20 +4,26 @@ * Runs a batch of agents on behalf of a caller agent: builds an * `AgentRunBatchLauncher` on top of the `agentLifecycle` primitives * (`create({ binding })`, `run`), drives the internal `AgentRunBatch` - * scheduler, and tracks one `AbortController` per caller so `cancel` can abort - * every in-flight run. The caller ↔ child association is this domain's own - * business data: requester-side display facts (`subagent.spawned` wire signals - * carrying the swarm's tool-call context, `subagent.suspended` when a task is + * scheduler, and tracks each caller's live batches so cancellation can target + * one member or all work owned by the caller. The caller ↔ child association + * is this domain's own business data: requester-side display facts + * (`subagent.spawned` wire signals carrying the swarm's tool-call context, + * `subagent.suspended` when a task is * requeued after a provider rate limit) are emitted here / via the * `agentLifecycle` wrapper helper `mirrorAgentRun`; the lifecycle registry - * itself stays flat. Bound at Session scope. + * itself stays flat. Repeated member stops remain idempotent without retaining + * every completed member for the whole session. Bound at Session scope. */ import type { TokenUsage } from '#/app/llmProtocol/usage'; import { InstantiationType } from '#/_base/di/extensions'; import { LifecycleScope, registerScopedService } from '#/_base/di/scope'; -import { linkAbortSignal } from '#/_base/utils/abort'; +import { + isUserCancellation, + linkAbortSignal, + userCancellationReason, +} from '#/_base/utils/abort'; import type { IAgentScopeHandle } from '#/_base/di/scope'; import { IAgentProfileService } from '#/agent/profile/profile'; import { IAgentPermissionModeService } from '#/agent/permissionMode/permissionMode'; @@ -44,6 +50,7 @@ import { ISessionSwarmService, type SessionSwarmRunArgs, type SessionSwarmRunResult, + type SessionSwarmStopResult, type SessionSwarmTask, } from './sessionSwarm'; import { @@ -66,11 +73,27 @@ declare module '#/app/event/eventBus' { * Kept as the legacy wire display value. */ const RESUMED_PROFILE_FALLBACK = 'subagent'; +const RECENT_SWARM_TERMINAL_LIMIT = 128; + +type InFlightBatch = { + readonly controller: AbortController; + readonly batch: Pick, 'stopAgent'>; + readonly agentIds: Set; +}; + +type CallerInFlight = { + readonly batches: Set; + readonly byAgentId: Map; +}; export class SessionSwarmService implements ISessionSwarmService { declare readonly _serviceBrand: undefined; - private readonly inFlight = new Map(); + private readonly inFlightByCaller = new Map(); + private readonly recentTerminalStatuses = new Map< + string, + SessionSwarmRunResult['status'] + >(); constructor( @IAgentLifecycleService private readonly lifecycle: IAgentLifecycleService, @@ -94,16 +117,18 @@ export class SessionSwarmService implements ISessionSwarmService { run(args: SessionSwarmRunArgs): Promise[]> { const { callerAgentId, tasks } = args; const controller = new AbortController(); - this.inFlight.set(callerAgentId, controller); + let inFlight: InFlightBatch; const unlinks: Array<() => void> = []; const linkedTasks: SessionSwarmTask[] = tasks.map((task) => { if (task.signal !== undefined) unlinks.push(linkAbortSignal(task.signal, controller)); return { ...task, signal: controller.signal }; }); const launcher: AgentRunBatchLauncher = { - spawn: (options) => this.spawnAttempt(callerAgentId, options), - resume: (agentId, options) => this.resumeAttempt(callerAgentId, agentId, options, false), - retry: (agentId, options) => this.resumeAttempt(callerAgentId, agentId, options, true), + spawn: (options) => this.spawnAttempt(callerAgentId, inFlight, options), + resume: (agentId, options) => + this.resumeAttempt(callerAgentId, inFlight, agentId, options, false), + retry: (agentId, options) => + this.resumeAttempt(callerAgentId, inFlight, agentId, options, true), suspended: (event) => { const caller = this.lifecycle.getHandle(callerAgentId); caller?.accessor.get(IEventBus)?.publish({ @@ -114,20 +139,70 @@ export class SessionSwarmService implements ISessionSwarmService { }, }; const maxConcurrency = resolveSwarmMaxConcurrency(); - const promise = new AgentRunBatch(launcher, linkedTasks, { maxConcurrency }).run(); - void promise.finally(() => { + const batch = new AgentRunBatch(launcher, linkedTasks, { maxConcurrency }); + inFlight = { controller, batch, agentIds: new Set() }; + this.addInFlight(callerAgentId, inFlight); + const promise = batch.run(); + const cleanup = () => { for (const unlink of unlinks) unlink(); - if (this.inFlight.get(callerAgentId) === controller) this.inFlight.delete(callerAgentId); - }); + this.removeInFlight(callerAgentId, inFlight); + }; + void promise.then( + (results) => { + for (const result of results) { + if (result.agentId === undefined) continue; + this.rememberTerminal(callerAgentId, result.agentId, result.status); + } + cleanup(); + }, + () => { + cleanup(); + }, + ); return promise; } + stopAgent(args: { + readonly callerAgentId: string; + readonly agentId: string; + }): SessionSwarmStopResult { + const caller = this.inFlightByCaller.get(args.callerAgentId); + const inFlight = caller?.byAgentId.get(args.agentId); + const current = inFlight?.batch.stopAgent(args.agentId); + if (current !== undefined && current.kind !== 'not_found') { + if (current.kind === 'stopped') { + const callerHandle = this.lifecycle.getHandle(args.callerAgentId); + if (callerHandle !== undefined) { + this.publishCancellation(callerHandle, args.agentId); + } + } + return current; + } + + const status = this.recentTerminalStatus( + this.agentRunKey(args.callerAgentId, args.agentId), + ); + if (status !== undefined) { + return { + kind: 'already_terminal', + agentId: args.agentId, + status, + }; + } + return { kind: 'not_found', agentId: args.agentId }; + } + cancel({ callerAgentId }: { readonly callerAgentId: string }): void { - this.inFlight.get(callerAgentId)?.abort(); + const caller = this.inFlightByCaller.get(callerAgentId); + if (caller === undefined) return; + for (const inFlight of caller.batches) { + inFlight.controller.abort(userCancellationReason()); + } } private async spawnAttempt( callerAgentId: string, + inFlight: InFlightBatch, options: AgentSpawnAttemptOptions, ): Promise { options.signal.throwIfAborted(); @@ -153,56 +228,88 @@ export class SessionSwarmService implements ISessionSwarmService { permissionMode: caller.accessor.get(IAgentPermissionModeService).mode, labels: subagentLabels(callerAgentId, { swarmItem: options.swarmItem }), }); + if (options.signal.aborted) { + await this.lifecycle.remove(child.id); + options.signal.throwIfAborted(); + } child.accessor .get(IAgentUserToolService) .inheritUserTools(caller.accessor.get(IAgentUserToolService)); - emitAgentRunSpawned(caller, child.id, { - profileName: options.profileName, - parentToolCallId: options.parentToolCallId, - parentToolCallUuid: options.parentToolCallUuid, - description: options.description, - swarmIndex: options.swarmIndex, - runInBackground: options.runInBackground, - }); const promptText = await applyProfilePromptPrefix(profile, options.prompt, { cwd: this.sessionContext.cwd, runner: this.processRunner, log: this.log, }); - return this.observe(caller, child.id, options.profileName, { - kind: 'prompt', - prompt: promptText, - }, options); + if (options.signal.aborted) { + await this.lifecycle.remove(child.id); + options.signal.throwIfAborted(); + } + let announced = false; + try { + this.identifyAgent(callerAgentId, inFlight, child.id, options); + emitAgentRunSpawned(caller, child.id, { + profileName: options.profileName, + parentToolCallId: options.parentToolCallId, + parentToolCallUuid: options.parentToolCallUuid, + description: options.description, + swarmIndex: options.swarmIndex, + runInBackground: options.runInBackground, + }); + announced = true; + return await this.observe(caller, child.id, options.profileName, { + kind: 'prompt', + prompt: promptText, + }, options); + } catch (error) { + if (announced && isUserCancellation(options.signal.reason)) { + this.publishCancellation(caller, child.id); + } + throw error; + } } private async resumeAttempt( callerAgentId: string, + inFlight: InFlightBatch, agentId: string, options: AgentRunAttemptOptions, retryTurn: boolean, ): Promise { options.signal.throwIfAborted(); - await this.requireOwnedSubagent(callerAgentId, agentId); const caller = this.requireHandle(callerAgentId, 'Caller agent'); - const child = this.requireHandle(agentId, 'Agent instance'); - this.requireIdleSubagent(agentId, child); - this.realignChildModel(caller, child); - const profileName = - child.accessor.get(IAgentProfileService).data().profileName ?? RESUMED_PROFILE_FALLBACK; - if (!retryTurn) { - emitAgentRunSpawned(caller, agentId, { - profileName, - parentToolCallId: options.parentToolCallId, - parentToolCallUuid: options.parentToolCallUuid, - description: options.description, - swarmIndex: options.swarmIndex, - runInBackground: options.runInBackground, - }); + let announced = retryTurn; + try { + await this.requireOwnedSubagent(callerAgentId, agentId); + options.signal.throwIfAborted(); + const child = this.requireHandle(agentId, 'Agent instance'); + this.requireIdleSubagent(agentId, child); + this.realignChildModel(caller, child); + const profileName = + child.accessor.get(IAgentProfileService).data().profileName ?? RESUMED_PROFILE_FALLBACK; + if (retryTurn) { + this.requireAgentOwner(callerAgentId, inFlight, agentId, options.signal); + } else { + this.identifyAgent(callerAgentId, inFlight, agentId, options); + emitAgentRunSpawned(caller, agentId, { + profileName, + parentToolCallId: options.parentToolCallId, + parentToolCallUuid: options.parentToolCallUuid, + description: options.description, + swarmIndex: options.swarmIndex, + runInBackground: options.runInBackground, + }); + announced = true; + } + const request = retryTurn + ? ({ kind: 'retry' } as const) + : ({ kind: 'prompt', prompt: options.prompt } as const); + return await this.observe(caller, child.id, profileName, request, options); + } catch (error) { + if (announced && isUserCancellation(options.signal.reason)) { + this.publishCancellation(caller, agentId); + } + throw error; } - const request = retryTurn - ? ({ kind: 'retry' } as const) - : ({ kind: 'prompt', prompt: options.prompt } as const); - return this.observe(caller, child.id, profileName, request, options); } private async observe( @@ -212,6 +319,7 @@ export class SessionSwarmService implements ISessionSwarmService { request: { kind: 'prompt'; prompt: string } | { kind: 'retry' }, options: AgentRunAttemptOptions, ): Promise { + options.signal.throwIfAborted(); const run = await this.lifecycle.run(agentId, request, { signal: options.signal, onReady: options.onReady, @@ -229,6 +337,98 @@ export class SessionSwarmService implements ISessionSwarmService { }; } + private identifyAgent( + callerAgentId: string, + inFlight: InFlightBatch, + agentId: string, + options: AgentRunAttemptOptions, + ): void { + const caller = this.inFlightByCaller.get(callerAgentId); + if (caller === undefined || !caller.batches.has(inFlight)) { + options.signal.throwIfAborted(); + throw new Error('Swarm batch is no longer running'); + } + if (caller.byAgentId.has(agentId)) { + throw new Error(`Agent instance "${agentId}" is already owned by a running swarm batch`); + } + this.recentTerminalStatuses.delete(this.agentRunKey(callerAgentId, agentId)); + inFlight.agentIds.add(agentId); + caller.byAgentId.set(agentId, inFlight); + options.onAgentIdentified?.(agentId); + } + + private publishCancellation( + caller: IAgentScopeHandle, + agentId: string, + ): void { + caller.accessor.get(IEventBus)?.publish({ + type: 'subagent.failed', + subagentId: agentId, + error: 'Aborted by the user', + cancelled: true, + }); + } + + private requireAgentOwner( + callerAgentId: string, + inFlight: InFlightBatch, + agentId: string, + signal: AbortSignal, + ): void { + const caller = this.inFlightByCaller.get(callerAgentId); + if (caller?.byAgentId.get(agentId) === inFlight) return; + signal.throwIfAborted(); + throw new Error(`Agent instance "${agentId}" is not owned by this swarm batch`); + } + + private addInFlight(callerAgentId: string, inFlight: InFlightBatch): void { + let caller = this.inFlightByCaller.get(callerAgentId); + if (caller === undefined) { + caller = { batches: new Set(), byAgentId: new Map() }; + this.inFlightByCaller.set(callerAgentId, caller); + } + caller.batches.add(inFlight); + } + + private removeInFlight(callerAgentId: string, inFlight: InFlightBatch): void { + const caller = this.inFlightByCaller.get(callerAgentId); + if (caller === undefined) return; + caller.batches.delete(inFlight); + for (const agentId of inFlight.agentIds) { + if (caller.byAgentId.get(agentId) === inFlight) caller.byAgentId.delete(agentId); + } + if (caller.batches.size === 0) this.inFlightByCaller.delete(callerAgentId); + } + + private rememberTerminal( + callerAgentId: string, + agentId: string, + status: SessionSwarmRunResult['status'], + ): void { + const key = this.agentRunKey(callerAgentId, agentId); + this.recentTerminalStatuses.delete(key); + this.recentTerminalStatuses.set(key, status); + while (this.recentTerminalStatuses.size > RECENT_SWARM_TERMINAL_LIMIT) { + const oldest = this.recentTerminalStatuses.keys().next().value; + if (oldest === undefined) break; + this.recentTerminalStatuses.delete(oldest); + } + } + + private recentTerminalStatus( + key: string, + ): SessionSwarmRunResult['status'] | undefined { + const status = this.recentTerminalStatuses.get(key); + if (status === undefined) return undefined; + this.recentTerminalStatuses.delete(key); + this.recentTerminalStatuses.set(key, status); + return status; + } + + private agentRunKey(callerAgentId: string, agentId: string): string { + return `${callerAgentId}\0${agentId}`; + } + private requireHandle(agentId: string, label: string): IAgentScopeHandle { const handle = this.lifecycle.getHandle(agentId); if (handle === undefined) throw new Error(`${label} "${agentId}" does not exist`); diff --git a/packages/agent-core-v2/test/agent/swarm/swarm.test.ts b/packages/agent-core-v2/test/agent/swarm/swarm.test.ts index 8d5a8637ee..bca2d1ea37 100644 --- a/packages/agent-core-v2/test/agent/swarm/swarm.test.ts +++ b/packages/agent-core-v2/test/agent/swarm/swarm.test.ts @@ -53,7 +53,13 @@ function mockSwarmHost({ readonly getSwarmItem?: (...args: any[]) => any; } = {}) { return { - swarmService: { _serviceBrand: undefined, getSwarmItem, run, cancel: vi.fn() }, + swarmService: { + _serviceBrand: undefined, + getSwarmItem, + run, + stopAgent: ({ agentId }: { agentId: string }) => ({ kind: 'not_found' as const, agentId }), + cancel: vi.fn(), + }, callerAgentId: 'main', }; } @@ -84,6 +90,7 @@ describe('AgentSwarmService', () => { ix.stub(ISessionSwarmService, { getSwarmItem: async () => undefined, run: async () => [], + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, }); ix.stub(IAgentScopeContext, makeAgentScopeContext({ agentId: 'main', agentScope: '' })); diff --git a/packages/agent-core-v2/test/agent/task/taskManager.test.ts b/packages/agent-core-v2/test/agent/task/taskManager.test.ts index 7ec4c5203c..c62cb332b9 100644 --- a/packages/agent-core-v2/test/agent/task/taskManager.test.ts +++ b/packages/agent-core-v2/test/agent/task/taskManager.test.ts @@ -513,6 +513,54 @@ describe('AgentTaskService', () => { expect(isUserCancellation(subagentController.signal.reason)).toBe(true); }); + it('stops an Agent task with a user-cancellation signal', async () => { + const { manager } = createAgentTaskService(); + const subagentController = new AbortController(); + const completion = new Promise<{ result: string }>((_resolve, reject) => { + subagentController.signal.addEventListener( + 'abort', + () => reject(subagentController.signal.reason), + { once: true }, + ); + }); + const taskId = manager.registerTask( + agentTask(completion, 'background agent', { abortController: subagentController }), + ); + + await manager.stop(taskId, 'Stopped from the web'); + + expect(isUserCancellation(subagentController.signal.reason)).toBe(true); + expect(manager.getTask(taskId)).toMatchObject({ + status: 'killed', + stopReason: 'Stopped from the web', + }); + }); + + it('resolves an Agent task by its stable Agent id after it stops', async () => { + const { manager } = createAgentTaskService(); + const subagentController = new AbortController(); + const completion = new Promise<{ result: string }>((_resolve, reject) => { + subagentController.signal.addEventListener( + 'abort', + () => reject(subagentController.signal.reason), + { once: true }, + ); + }); + const taskId = manager.registerTask( + agentTask(completion, 'background agent', { + agentId: 'agent-stable', + abortController: subagentController, + }), + ); + + expect(manager.getAgentTask('agent-stable')?.taskId).toBe(taskId); + await manager.stop(taskId); + expect(manager.getAgentTask('agent-stable')).toMatchObject({ + taskId, + status: 'killed', + }); + }); + it('does not count foreground tasks against the detached task limit', () => { const { manager } = createAgentTaskService({ maxRunningTasks: 1 }); manager.registerTask(agentTask(new Promise(() => {}), 'foreground agent'), { diff --git a/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts b/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts index 9b8e1e4e43..ee5ca349a1 100644 --- a/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts +++ b/packages/agent-core-v2/test/agent/task/tools/task-tools.test.ts @@ -158,6 +158,12 @@ class FakeTaskService implements IAgentTaskService { return this.entries.get(taskId)?.info; } + getAgentTask(agentId: string): AgentTaskInfo | undefined { + return Array.from(this.entries.values(), (entry) => entry.info) + .filter((info) => info.kind === 'agent' && info.agentId === agentId) + .toSorted((left, right) => right.startedAt - left.startedAt)[0]; + } + list(activeOnly = true, limit?: number): readonly AgentTaskInfo[] { const result: AgentTaskInfo[] = []; for (const entry of this.entries.values()) { diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index d88b1feb9c..fab49f0886 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -654,6 +654,7 @@ export function swarmServices( _serviceBrand: undefined, getSwarmItem: async () => undefined, run: swarmService, + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, } satisfies ISessionSwarmService : swarmService; diff --git a/packages/agent-core-v2/test/os/backends/node-local/hostProcessService.test.ts b/packages/agent-core-v2/test/os/backends/node-local/hostProcessService.test.ts index a84a2c00f1..18bd1dad3d 100644 --- a/packages/agent-core-v2/test/os/backends/node-local/hostProcessService.test.ts +++ b/packages/agent-core-v2/test/os/backends/node-local/hostProcessService.test.ts @@ -1,5 +1,13 @@ -import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +/** + * Scenario: node-local process spawning and tree-scoped termination. + * Exercises `IHostProcessService` through DI with real local processes. + * Run: `pnpm exec vitest run test/os/backends/node-local/hostProcessService.test.ts`. + */ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { once } from 'node:events'; +import { createInterface, type Interface as ReadlineInterface } from 'node:readline'; import { Readable } from 'node:stream'; import { DisposableStore } from '#/_base/di/lifecycle'; @@ -8,9 +16,30 @@ import { HostProcessError, HostProcessErrorCode, IHostProcessService, + type IHostProcess, } from '#/os/interface/hostProcess'; import { HostProcessService } from '#/os/backends/node-local/hostProcessService'; +const PROCESS_TREE_SCRIPT = ` + const { spawn } = require('node:child_process'); + const readline = require('node:readline'); + + const child = spawn(process.execPath, ['-e', 'setInterval(() => {}, 1000)'], { + stdio: 'ignore', + }); + process.stdout.write('READY:' + child.pid + '\\n'); + readline.createInterface({ input: process.stdin }).on('line', (line) => { + if (line === 'PING') process.stdout.write('PONG\\n'); + }); + setInterval(() => {}, 1000); +`; + +interface ProcessTree { + readonly process: IHostProcess; + readonly lines: ReadlineInterface; + readonly grandchildPid: number; +} + async function collect(stream: Readable): Promise { const chunks: Buffer[] = []; for await (const chunk of stream) { @@ -19,6 +48,75 @@ async function collect(stream: Readable): Promise { return Buffer.concat(chunks).toString('utf8'); } +async function nextLine(lines: ReadlineInterface): Promise { + const [line] = await once(lines, 'line', { signal: AbortSignal.timeout(5_000) }); + return String(line); +} + +async function spawnProcessTree(service: IHostProcessService): Promise { + const proc = await service.spawn(process.execPath, ['-e', PROCESS_TREE_SCRIPT]); + const lines = createInterface({ input: proc.stdout, crlfDelay: Infinity }); + try { + const ready = await nextLine(lines); + const match = /^READY:(\d+)$/.exec(ready); + if (match?.[1] === undefined) { + throw new Error(`Process tree did not report a grandchild pid: ${ready}`); + } + return { process: proc, lines, grandchildPid: Number.parseInt(match[1], 10) }; + } catch (error) { + lines.close(); + await stopProcess(proc); + throw error; + } +} + +function isProcessAlive(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code === 'ESRCH') return false; + if (code === 'EPERM') return true; + throw error; + } +} + +async function waitForProcessExit(pid: number): Promise { + await vi.waitFor(() => expect(isProcessAlive(pid)).toBe(false), { + interval: 25, + timeout: 5_000, + }); +} + +async function stopProcess(proc: IHostProcess): Promise { + if (proc.exitCode === null) { + try { + await proc.kill('SIGKILL'); + } catch { + // Best-effort test cleanup; wait() below observes the real outcome. + } + } + try { + await proc.wait(); + } catch { + // Best-effort test cleanup. + } + proc.dispose(); +} + +async function cleanupProcessTree(tree: ProcessTree): Promise { + tree.lines.close(); + await stopProcess(tree.process); + if (!isProcessAlive(tree.grandchildPid)) return; + try { + process.kill(tree.grandchildPid, 'SIGKILL'); + } catch { + // Best-effort fallback if the process exited between the probe and kill. + } + await waitForProcessExit(tree.grandchildPid); +} + describe('HostProcessService', () => { let disposables: DisposableStore; let ix: TestInstantiationService; @@ -79,4 +177,27 @@ describe('HostProcessService', () => { const code = await proc.wait(); expect(code).not.toBe(0); }); + + it('kill() terminates only the selected process tree', async () => { + const svc = ix.get(IHostProcessService); + const trees: ProcessTree[] = []; + try { + const selected = await spawnProcessTree(svc); + trees.push(selected); + const survivor = await spawnProcessTree(svc); + trees.push(survivor); + + await selected.process.kill(); + await selected.process.wait(); + await waitForProcessExit(selected.grandchildPid); + + const pong = nextLine(survivor.lines); + survivor.process.stdin.write('PING\n'); + await expect(pong).resolves.toBe('PONG'); + expect(isProcessAlive(survivor.process.pid)).toBe(true); + expect(isProcessAlive(survivor.grandchildPid)).toBe(true); + } finally { + await Promise.all(trees.map(cleanupProcessTree)); + } + }, 15_000); }); diff --git a/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts b/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts index 6366de7bf4..906c85d415 100644 --- a/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts +++ b/packages/agent-core-v2/test/os/backends/node-local/tools/bash.test.ts @@ -571,6 +571,12 @@ function createFakeTaskService(options: { maxRunningTasks?: number } = {}): { return entry === undefined ? undefined : entryToInfo(entry); }, + getAgentTask(agentId: string): AgentTaskInfo | undefined { + return Array.from(tasks.values(), entryToInfo) + .filter((info) => info.kind === 'agent' && info.agentId === agentId) + .toSorted((left, right) => right.startedAt - left.startedAt)[0]; + }, + list(activeOnly = true): readonly AgentTaskInfo[] { const result: AgentTaskInfo[] = []; for (const entry of tasks.values()) { diff --git a/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts b/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts index 43c6488f23..2c2f6cee72 100644 --- a/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts +++ b/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts @@ -42,7 +42,11 @@ import { type AgentSpawnAttemptOptions, type QueuedAgentRunTask, } from '#/session/swarm/agentRunBatch'; -import { ISessionSwarmService, type SessionSwarmTask } from '#/session/swarm/sessionSwarm'; +import { + ISessionSwarmService, + type SessionSwarmStopResult, + type SessionSwarmTask, +} from '#/session/swarm/sessionSwarm'; import { SessionSwarmService } from '#/session/swarm/sessionSwarmService'; import { stubLog } from '../../_base/log/stubs'; @@ -212,6 +216,144 @@ describe('AgentRunBatch scheduling contract', () => { } }); + it('stops one active member without cancelling its sibling', async () => { + vi.useFakeTimers(); + try { + const { runBatch, attempts, stopAgent } = createMockAgentRunBatchRunner(); + const running = runBatch([queuedAgentRunTask(1), queuedAgentRunTask(2)]); + + await vi.advanceTimersByTimeAsync(0); + expect(attempts).toHaveLength(2); + expect(stopAgent('agent-1')).toEqual({ + kind: 'stopping', + agentId: 'agent-1', + }); + + attempts[1]!.outcome.resolve({ + task: attempts[1]!.task, + agentId: 'agent-2', + status: 'completed', + result: 'sibling completed', + }); + await vi.advanceTimersByTimeAsync(0); + + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-1', status: 'aborted', state: 'started' }, + { agentId: 'agent-2', status: 'completed', result: 'sibling completed' }, + ]); + } finally { + vi.useRealTimers(); + } + }); + + it('stops a suspended member without retrying it and closes the parent batch', async () => { + vi.useFakeTimers(); + try { + const suspended: AgentRunSuspendedEvent[] = []; + const { runBatch, attempts, stopAgent } = createMockAgentRunBatchRunner({ + onSuspended: (event) => suspended.push(event), + }); + const running = runBatch([queuedAgentRunTask(1), queuedAgentRunTask(2)]); + + await vi.advanceTimersByTimeAsync(0); + attempts[0]!.outcome.resolve({ type: 'rate_limited', agentId: 'agent-1' }); + await vi.advanceTimersByTimeAsync(0); + expect(suspended).toHaveLength(1); + + expect(stopAgent('agent-1')).toEqual({ + kind: 'stopped', + agentId: 'agent-1', + }); + await vi.advanceTimersByTimeAsync(3_000); + expect(attempts).toHaveLength(2); + + attempts[1]!.outcome.resolve({ + task: attempts[1]!.task, + agentId: 'agent-2', + status: 'completed', + result: 'sibling completed', + }); + await vi.advanceTimersByTimeAsync(0); + + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-1', status: 'aborted', state: 'started' }, + { agentId: 'agent-2', status: 'completed' }, + ]); + } finally { + vi.useRealTimers(); + } + }); + + it('lets a targeted stop win a provider-rate-limit race', async () => { + vi.useFakeTimers(); + try { + const suspended: AgentRunSuspendedEvent[] = []; + const { runBatch, attempts, stopAgent } = createMockAgentRunBatchRunner({ + onSuspended: (event) => suspended.push(event), + }); + const running = runBatch([queuedAgentRunTask(1), queuedAgentRunTask(2)]); + + await vi.advanceTimersByTimeAsync(0); + attempts[0]!.outcome.reject(new APIProviderRateLimitError('Rate limited')); + // Settle the attempt completion with the 429, but stop the member before + // AgentRunBatch handles that rejection. + await Promise.resolve(); + expect(stopAgent('agent-1').kind).toBe('stopping'); + await vi.advanceTimersByTimeAsync(0); + + attempts[1]!.outcome.resolve({ + task: attempts[1]!.task, + agentId: 'agent-2', + status: 'completed', + result: 'sibling completed', + }); + await vi.advanceTimersByTimeAsync(3_000); + + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-1', status: 'aborted' }, + { agentId: 'agent-2', status: 'completed' }, + ]); + expect(suspended).toEqual([]); + expect(attempts).toHaveLength(2); + } finally { + vi.useRealTimers(); + } + }); + + it('can stop a spawned member before its launcher returns the run handle', async () => { + const identified = createControlledPromise(); + const releaseSpawn = createControlledPromise(); + const launcher: AgentRunBatchLauncher = { + spawn: async (options) => { + options.onAgentIdentified?.('agent-race'); + identified.resolve(); + await releaseSpawn; + options.signal.throwIfAborted(); + return { + agentId: 'agent-race', + profileName: options.profileName, + completion: Promise.resolve({ result: 'unexpected' }), + }; + }, + resume: async () => { + throw new Error('unexpected resume'); + }, + retry: async () => { + throw new Error('unexpected retry'); + }, + }; + const batch = new AgentRunBatch(launcher, [queuedAgentRunTask(1)]); + const running = batch.run(); + + await identified; + expect(batch.stopAgent('agent-race').kind).toBe('stopping'); + releaseSpawn.resolve(); + + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-race', status: 'aborted', state: 'started' }, + ]); + }); + it('normal phase keeps processing completions while waiting for the next launch', async () => { vi.useFakeTimers(); try { @@ -903,6 +1045,41 @@ describe('SessionSwarmService metadata compatibility', () => { disposables.dispose(); }); + function registerResumableAgent(agentId: string): void { + agents[agentId] = { + homedir: `/tmp/kimi/s1/agents/${agentId}`, + labels: { parentAgentId: 'main' }, + }; + handles.set(agentId, agentHandle(agentId, lifecycle, eventBus)); + } + + function controlledResumableRuns(...agentIds: string[]) { + const completions = new Map( + agentIds.map( + (agentId) => + [agentId, createControlledPromise<{ summary: string }>()] as const, + ), + ); + const signals = new Map(); + const started = createControlledPromise(); + agentIds.forEach(registerResumableAgent); + runAgent.mockImplementation((agentId, _request, options) => { + const signal = options?.signal; + const completion = completions.get(agentId); + if (signal === undefined || completion === undefined) { + throw new Error(`Unexpected swarm run for ${agentId}`); + } + signals.set(agentId, signal); + signal.addEventListener('abort', () => completion.reject(signal.reason), { + once: true, + }); + options.onReady?.(); + if (signals.size === agentIds.length) started.resolve(); + return { agentId, turn: {} as never, completion }; + }); + return { completions, signals, started }; + } + it('reads swarm items from caller-owned v2 labels and legacy v1 metadata', async () => { agents['v2-child'] = { homedir: '/tmp/kimi/s1/agents/v2-child', @@ -1153,6 +1330,303 @@ describe('SessionSwarmService metadata compatibility', () => { } }); + it('publishes a cancelled lifecycle event when a suspended member is stopped', async () => { + vi.useFakeTimers(); + try { + ['agent-rate-limited', 'agent-sibling'].forEach(registerResumableAgent); + const rateLimited = createControlledPromise<{ summary: string }>(); + const sibling = createControlledPromise<{ summary: string }>(); + const published: DomainEvent[] = []; + (eventBus.publish as ReturnType).mockImplementation((event: DomainEvent) => { + published.push(event); + }); + runAgent.mockImplementation((agentId, _request, options) => { + options?.onReady?.(); + return { + agentId, + turn: {} as never, + completion: agentId === 'agent-rate-limited' ? rateLimited : sibling, + }; + }); + const service = ix.get(ISessionSwarmService); + const running = service.run({ + callerAgentId: 'main', + tasks: [ + resumeSessionTask('agent-rate-limited'), + resumeSessionTask('agent-sibling'), + ], + }); + await vi.advanceTimersByTimeAsync(0); + rateLimited.reject(new APIProviderRateLimitError('Rate limited')); + await vi.advanceTimersByTimeAsync(0); + + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-rate-limited' }), + ).toMatchObject({ kind: 'stopped' }); + expect(published).toContainEqual( + expect.objectContaining({ + type: 'subagent.failed', + subagentId: 'agent-rate-limited', + cancelled: true, + }), + ); + + sibling.resolve({ summary: 'sibling summary' }); + await vi.advanceTimersByTimeAsync(0); + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-rate-limited', status: 'aborted' }, + { agentId: 'agent-sibling', status: 'completed' }, + ]); + } finally { + vi.useRealTimers(); + } + }); + + it('keeps a stopped member terminal after its batch has been cleaned up', async () => { + const completion = createControlledPromise<{ summary: string }>(); + runAgent.mockImplementation((agentId, _request, options) => { + options?.onReady?.(); + options?.signal.addEventListener( + 'abort', + () => completion.reject(options.signal.reason), + { once: true }, + ); + return { agentId, turn: {} as never, completion }; + }); + const service = ix.get(ISessionSwarmService); + const running = service.run({ + callerAgentId: 'main', + tasks: [spawnSessionTask('src/a.ts')], + }); + await vi.waitFor(() => expect(runAgent).toHaveBeenCalledOnce()); + + expect(service.stopAgent({ callerAgentId: 'main', agentId: 'agent-new' })).toEqual({ + kind: 'stopping', + agentId: 'agent-new', + }); + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-new', status: 'aborted', state: 'started' }, + ]); + + expect(service.stopAgent({ callerAgentId: 'main', agentId: 'agent-new' })).toEqual({ + kind: 'already_terminal', + agentId: 'agent-new', + status: 'aborted', + }); + expect(lifecycle.getHandle('agent-new')).toBeDefined(); + }); + + it('stops only the targeted member when one caller has concurrent batches', async () => { + const runs = controlledResumableRuns('agent-first-batch', 'agent-second-batch'); + const service = ix.get(ISessionSwarmService); + const firstRunning = service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-first-batch')], + }); + const secondRunning = service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-second-batch')], + }); + await runs.started; + + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-first-batch' }), + ).toEqual({ + kind: 'stopping', + agentId: 'agent-first-batch', + }); + expect(runs.signals.get('agent-second-batch')?.aborted).toBe(false); + runs.completions + .get('agent-second-batch')! + .resolve({ summary: 'second batch completed' }); + + await expect(firstRunning).resolves.toMatchObject([ + { agentId: 'agent-first-batch', status: 'aborted' }, + ]); + await expect(secondRunning).resolves.toMatchObject([ + { + agentId: 'agent-second-batch', + status: 'completed', + result: 'second batch completed', + }, + ]); + }); + + it('keeps the first batch ownership when the same agent is resumed concurrently', async () => { + const runs = controlledResumableRuns('agent-shared'); + const service = ix.get(ISessionSwarmService); + const firstRunning = service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-shared')], + }); + await runs.started; + + await expect( + service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-shared')], + }), + ).resolves.toMatchObject([ + { + status: 'failed', + state: 'not_started', + error: 'Agent instance "agent-shared" is already owned by a running swarm batch', + }, + ]); + expect(runAgent).toHaveBeenCalledOnce(); + expect(service.stopAgent({ callerAgentId: 'main', agentId: 'agent-shared' })).toEqual({ + kind: 'stopping', + agentId: 'agent-shared', + }); + await expect(firstRunning).resolves.toMatchObject([ + { agentId: 'agent-shared', status: 'aborted' }, + ]); + }); + + it('cancels every concurrent batch owned by one caller', async () => { + const runs = controlledResumableRuns('agent-cancel-first', 'agent-cancel-second'); + const service = ix.get(ISessionSwarmService); + const firstRunning = service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-cancel-first')], + }); + const secondRunning = service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-cancel-second')], + }); + await runs.started; + + service.cancel({ callerAgentId: 'main' }); + + expect(runs.signals.get('agent-cancel-first')?.aborted).toBe(true); + expect(runs.signals.get('agent-cancel-second')?.aborted).toBe(true); + await expect(firstRunning).resolves.toMatchObject([ + { agentId: 'agent-cancel-first', status: 'aborted' }, + ]); + await expect(secondRunning).resolves.toMatchObject([ + { agentId: 'agent-cancel-second', status: 'aborted' }, + ]); + }); + + it('removes a child created after its batch was cancelled before emitting spawned', async () => { + const createStarted = createControlledPromise(); + const releaseCreate = createControlledPromise(); + createAgent.mockImplementationOnce(async () => { + createStarted.resolve(); + await releaseCreate; + const child = agentHandle('agent-created-after-cancel', lifecycle, eventBus); + handles.set(child.id, child); + return child; + }); + const service = ix.get(ISessionSwarmService); + const running = service.run({ + callerAgentId: 'main', + tasks: [spawnSessionTask('src/cancelled.ts')], + }); + await createStarted; + + service.cancel({ callerAgentId: 'main' }); + releaseCreate.resolve(); + + await expect(running).resolves.toMatchObject([{ status: 'aborted' }]); + expect(lifecycle.getHandle('agent-created-after-cancel')).toBeUndefined(); + expect(runAgent).not.toHaveBeenCalled(); + expect(eventBus.publish).not.toHaveBeenCalledWith( + expect.objectContaining({ + type: 'subagent.spawned', + subagentId: 'agent-created-after-cancel', + }), + ); + }); + + it('forgets the oldest terminal member after retaining 128 recent members', async () => { + for (let index = 0; index < 129; index += 1) { + registerResumableAgent(`agent-terminal-${String(index)}`); + } + const service = ix.get(ISessionSwarmService); + + for (let index = 0; index < 129; index += 1) { + await service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask(`agent-terminal-${String(index)}`)], + }); + } + + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-terminal-0' }), + ).toEqual({ kind: 'not_found', agentId: 'agent-terminal-0' }); + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-terminal-128' }), + ).toEqual({ + kind: 'already_terminal', + agentId: 'agent-terminal-128', + status: 'completed', + }); + }); + + it('retains a recently accessed terminal member when the LRU reaches capacity', async () => { + for (let index = 0; index < 129; index += 1) { + registerResumableAgent(`agent-lru-${String(index)}`); + } + const service = ix.get(ISessionSwarmService); + for (let index = 0; index < 128; index += 1) { + await service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask(`agent-lru-${String(index)}`)], + }); + } + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-lru-0' }), + ).toEqual({ + kind: 'already_terminal', + agentId: 'agent-lru-0', + status: 'completed', + }); + + await service.run({ + callerAgentId: 'main', + tasks: [resumeSessionTask('agent-lru-128')], + }); + + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-lru-1' }), + ).toEqual({ kind: 'not_found', agentId: 'agent-lru-1' }); + expect( + service.stopAgent({ callerAgentId: 'main', agentId: 'agent-lru-0' }), + ).toEqual({ + kind: 'already_terminal', + agentId: 'agent-lru-0', + status: 'completed', + }); + }); + + it('does not stop an in-flight member owned by another caller', async () => { + const completion = createControlledPromise<{ summary: string }>(); + let runSignal: AbortSignal | undefined; + runAgent.mockImplementation((agentId, _request, options) => { + runSignal = options?.signal; + options?.onReady?.(); + return { agentId, turn: {} as never, completion }; + }); + const service = ix.get(ISessionSwarmService); + const running = service.run({ + callerAgentId: 'main', + tasks: [spawnSessionTask('src/a.ts')], + }); + await vi.waitFor(() => expect(runAgent).toHaveBeenCalledOnce()); + + expect(service.stopAgent({ callerAgentId: 'other', agentId: 'agent-new' })).toEqual({ + kind: 'not_found', + agentId: 'agent-new', + }); + expect(runSignal?.aborted).toBe(false); + + completion.resolve({ summary: 'child summary' }); + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-new', status: 'completed' }, + ]); + }); + it('rejects resume of an already running child before launching or emitting spawned', async () => { agents['agent-existing'] = { homedir: '/tmp/kimi/s1/agents/agent-existing', @@ -1367,9 +1841,11 @@ function createMockAgentRunBatchRunner( options?: { readonly signal?: AbortSignal }, ) => Promise>>; readonly attempts: MockAgentRunAttemptRecord[]; + readonly stopAgent: (agentId: string) => SessionSwarmStopResult; } { const attempts: MockAgentRunAttemptRecord[] = []; let activeTasks: readonly QueuedAgentRunTask[] = []; + let activeBatch: AgentRunBatch | undefined; const createHandle = ( runOptions: AgentRunAttemptOptions, @@ -1425,11 +1901,15 @@ function createMockAgentRunBatchRunner( ...task, signal: task.signal ?? runOptions?.signal, })); - return new AgentRunBatch(launcher, activeTasks as readonly QueuedAgentRunTask[], { + const batch = new AgentRunBatch(launcher, activeTasks as readonly QueuedAgentRunTask[], { maxConcurrency: options.maxConcurrency, - }).run(); + }); + activeBatch = batch as AgentRunBatch; + return batch.run(); }, attempts, + stopAgent: (agentId) => + activeBatch?.stopAgent(agentId) ?? { kind: 'not_found', agentId }, }; } diff --git a/packages/agent-core-v2/test/tool/tool.test.ts b/packages/agent-core-v2/test/tool/tool.test.ts index fa73bb7e84..bf29006225 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -650,6 +650,50 @@ describe('Agent tool execution contract', () => { }); }); + it('mirrors explicit user cancellation as a cancelled subagent failure', async () => { + const lifecycle = createAgentLifecycleStub(); + const events: DomainEvent[] = []; + const eventBus = { + _serviceBrand: undefined, + publish: vi.fn((event: DomainEvent) => events.push(event)), + subscribe: vi.fn(() => noopDisposable()), + } as IEventBus; + const requester = { + id: 'main', + kind: LifecycleScope.Agent, + accessor: { + get: ((serviceId: unknown) => { + if (serviceId === IEventBus) return eventBus; + if (serviceId === IAgentLifecycleService) return lifecycle; + return undefined; + }) as IAgentScopeHandle['accessor']['get'], + }, + dispose: () => {}, + } satisfies IAgentScopeHandle; + const controller = new AbortController(); + const reason = userCancellationReason(); + controller.abort(reason); + + const mirrored = mirrorAgentRun( + requester, + { + agentId: 'agent-child', + turn: {} as AgentRunHandle['turn'], + completion: Promise.reject(reason), + }, + { + profileName: 'explore', + signal: controller.signal, + }, + ); + + await expect(mirrored).rejects.toBe(reason); + expect(events.find((event) => event.type === 'subagent.failed')).toMatchObject({ + subagentId: 'agent-child', + cancelled: true, + }); + }); + it('inherits parent user tools when spawning a subagent', async () => { const lookupTool: UserToolRegistration = { name: 'Lookup', @@ -1409,6 +1453,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1490,6 +1535,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem, run: runSwarm as ISessionSwarmService['run'], + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1615,6 +1661,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1662,6 +1709,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1718,6 +1766,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], + stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); diff --git a/packages/kap-server/src/routes/tasks.ts b/packages/kap-server/src/routes/tasks.ts index 922367a11c..4f9f7b354e 100644 --- a/packages/kap-server/src/routes/tasks.ts +++ b/packages/kap-server/src/routes/tasks.ts @@ -9,13 +9,16 @@ * output_bytes?} data: Task * POST /sessions/{session_id}/tasks/{task_id}:cancel body: empty data: {cancelled:true} * - * **Thin wrapper over `IAgentTaskService`**: the main agent's + * **Thin wrapper over existing task domains**: the main agent's * `IAgentTaskService` is already exposed at Agent scope (`tasks:*` in the RPC * action map). These REST routes borrow it by interface and project its * `AgentTaskInfo` (camelCase + ms timestamps + agent-core literal sets) - * into the protocol's `Task` shape (snake_case + ISO + spec literal - * sets) — the same field/literal mapping v1 performs in - * `packages/agent-core/src/services/task/task.ts`. + * into the protocol's `Task` shape (snake_case + ISO + spec literal sets). + * Cancel accepts Agent ids as compatibility aliases for Agent tasks. Swarm + * members are not task-service entries, so cancellation falls through to the + * Session-scoped `ISessionSwarmService`, which owns member isolation. Cancel + * resolution prefers an exact task id, then an active Agent-task alias, then a + * swarm member, and only then terminal Agent-task history. * * **Resolution**: `core` → `ISessionIndex` (existence, → 40401) → * `ISessionLifecycleService` (live session handle) → `IAgentLifecycleService` @@ -42,7 +45,9 @@ import { IAgentTaskService, ISessionIndex, ISessionLifecycleService, + ISessionSwarmService, type AgentTaskInfo, + type SessionSwarmStopResult, type Scope, } from '@moonshot-ai/agent-core-v2'; import { @@ -58,7 +63,7 @@ import { z } from 'zod'; import { errEnvelope, okEnvelope } from '../envelope'; import { defineRoute } from '../middleware/defineRoute'; -import { ensureMainAgent } from '../transport/mainAgent'; +import { ensureMainAgent, MAIN_AGENT_ID } from '../transport/mainAgent'; import { parseActionSuffix } from './action-suffix'; /** Default cap (bytes) for the opt-in output preview on GET-by-id. */ @@ -241,22 +246,39 @@ export function registerTasksRoutes(app: TasksRouteHost, core: Scope): void { return; } - // Pre-fetch so we can distinguish 40406 (not found) from 40904 (already - // finished) deterministically — `IAgentTaskService.stop` does not - // surface this distinction on its own. - const found = resolved.tasks?.getTask(task_id); - if (found === undefined) { - reply.send(taskNotFound(session_id, task_id, req.id)); + const target = resolveCancelTarget(resolved, task_id); + if (target.kind === 'agent_task') { + const wireStatus = toWireTask(session_id, target.task).status; + if (isTerminalStatus(wireStatus)) { + reply.send(taskAlreadyFinished(session_id, task_id, wireStatus, req.id)); + return; + } + + await resolved.tasks?.stop(target.task.taskId); + reply.send(okEnvelope({ cancelled: true as const }, req.id)); return; } - const wireStatus = toWireTask(session_id, found).status; - if (isTerminalStatus(wireStatus)) { - reply.send(taskAlreadyFinished(session_id, task_id, wireStatus, req.id)); + + if ( + target.kind === 'swarm' && + (target.result.kind === 'stopping' || target.result.kind === 'stopped') + ) { + reply.send(okEnvelope({ cancelled: true as const }, req.id)); + return; + } + if (target.kind === 'swarm' && target.result.kind === 'already_terminal') { + reply.send( + taskAlreadyFinished( + session_id, + task_id, + mapSwarmStatus(target.result.status), + req.id, + ), + ); return; } - await resolved.tasks?.stop(task_id); - reply.send(okEnvelope({ cancelled: true as const }, req.id)); + reply.send(taskNotFound(session_id, task_id, req.id)); }, ); app.post(cancelRoute.path, cancelRoute.options, cancelRoute.handler as Parameters[2]); @@ -271,17 +293,56 @@ export function registerTasksRoutes(app: TasksRouteHost, core: Scope): void { type ResolvedTasks = | { readonly kind: 'not_found' } - | { readonly kind: 'resolved'; readonly tasks: IAgentTaskService | undefined }; + | { + readonly kind: 'resolved'; + readonly tasks: IAgentTaskService | undefined; + readonly swarm: ISessionSwarmService | undefined; + }; async function resolveSessionTasks(core: Scope, sid: string): Promise { const summary = await core.accessor.get(ISessionIndex).get(sid); if (summary === undefined) return { kind: 'not_found' }; const session = core.accessor.get(ISessionLifecycleService).get(sid); - if (session === undefined) return { kind: 'resolved', tasks: undefined }; + if (session === undefined) { + return { kind: 'resolved', tasks: undefined, swarm: undefined }; + } const agent = await ensureMainAgent(session); const tasks = agent.accessor.get(IAgentTaskService); - return { kind: 'resolved', tasks }; + const swarm = session.accessor.get(ISessionSwarmService); + return { kind: 'resolved', tasks, swarm }; +} + +type SwarmCancelResult = Exclude; + +type CancelTarget = + | { readonly kind: 'agent_task'; readonly task: AgentTaskInfo } + | { readonly kind: 'swarm'; readonly result: SwarmCancelResult } + | { readonly kind: 'not_found' }; + +function resolveCancelTarget( + resolved: Extract, + taskOrAgentId: string, +): CancelTarget { + const exactTask = resolved.tasks?.getTask(taskOrAgentId); + if (exactTask !== undefined) return { kind: 'agent_task', task: exactTask }; + + const agentTask = resolved.tasks?.getAgentTask(taskOrAgentId); + if (agentTask !== undefined && !isTerminalStatus(mapStatus(agentTask.status))) { + return { kind: 'agent_task', task: agentTask }; + } + + const swarmResult = resolved.swarm?.stopAgent({ + callerAgentId: MAIN_AGENT_ID, + agentId: taskOrAgentId, + }); + if (swarmResult !== undefined && swarmResult.kind !== 'not_found') { + return { kind: 'swarm', result: swarmResult }; + } + + return agentTask === undefined + ? { kind: 'not_found' } + : { kind: 'agent_task', task: agentTask }; } // --------------------------------------------------------------------------- @@ -333,6 +394,10 @@ function mapStatus(s: AgentTaskInfo['status']): TaskStatus { } } +function mapSwarmStatus(status: 'completed' | 'failed' | 'aborted'): TaskStatus { + return status === 'aborted' ? 'cancelled' : status; +} + const TERMINAL_WIRE_STATUSES: ReadonlySet = new Set([ 'completed', 'failed', diff --git a/packages/kap-server/test/tasks.test.ts b/packages/kap-server/test/tasks.test.ts index e59f58f18d..4487b8366c 100644 --- a/packages/kap-server/test/tasks.test.ts +++ b/packages/kap-server/test/tasks.test.ts @@ -6,10 +6,11 @@ import { IAgentLifecycleService, IAgentTaskService, ISessionLifecycleService, + ISessionSwarmService, IModelResolver, type AgentTask, } from '@moonshot-ai/agent-core-v2'; -import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { type RunningServer, startServer } from '../src/start'; import { authHeaders } from './helpers/auth'; @@ -67,6 +68,7 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { }); afterEach(async () => { + vi.restoreAllMocks(); if (server !== undefined) { await server.close(); server = undefined; @@ -103,6 +105,15 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { return body.data.id; } + function cancelTask( + sessionId: string, + taskOrAgentId: string, + ): Promise<{ status: number; body: Envelope }> { + return postJson( + `/api/v1/sessions/${sessionId}/tasks/${encodeURIComponent(taskOrAgentId)}:cancel`, + ); + } + // The main agent scope is not created automatically on session creation // (server-v2 gap G10); create it here, then register fake tasks // directly into its IAgentTaskService to bypass the tool loop. @@ -115,13 +126,24 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { return agent.accessor.get(IAgentTaskService); } + async function sessionSwarm(sessionId: string): Promise { + await mainAgentTasks(sessionId); + const session = server!.core.accessor.get(ISessionLifecycleService).get(sessionId); + if (session === undefined) throw new Error(`session ${sessionId} not found`); + return session.accessor.get(ISessionSwarmService); + } + // Let the `registerTask` microtask run `start` (which appends output) before // the next request. async function flush(): Promise { await new Promise((resolve) => setTimeout(resolve, 10)); } - function fakeTask(kind: 'process' | 'agent' | 'question', output?: string): AgentTask { + function fakeTask( + kind: 'process' | 'agent' | 'question', + output?: string, + agentId = 'sub-1', + ): AgentTask { return { idPrefix: 'test', kind, @@ -134,7 +156,7 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { case 'process': return { ...base, kind: 'process', command: 'echo hi', pid: 0, exitCode: null }; case 'agent': - return { ...base, kind: 'agent', agentId: 'sub-1', subagentType: 'explore' }; + return { ...base, kind: 'agent', agentId, subagentType: 'explore' }; case 'question': return { ...base, kind: 'question', questionCount: 1 }; } @@ -142,6 +164,15 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { }; } + function completedAgentTask(agentId: string): AgentTask { + return { + ...fakeTask('agent', undefined, agentId), + start: async (sink) => { + await sink.settle({ status: 'completed' }); + }, + }; + } + it('returns an empty list when the session has no main agent (gap G10)', async () => { const id = await createSession(); const { body } = await getJson(`/api/v1/sessions/${id}/tasks`); @@ -227,6 +258,32 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { expect(missing.body.code).toBe(40406); }); + it('cancels an Agent task by its stable Agent id', async () => { + const id = await createSession(); + const tasks = await mainAgentTasks(id); + const taskId = tasks.registerTask(fakeTask('agent', undefined, 'agent-ui')); + await flush(); + + const cancelled = await cancelTask(id, 'agent-ui'); + expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); + expect(tasks.getTask(taskId)).toMatchObject({ status: 'killed' }); + }); + + it('reports an Agent task as already finished when its alias is cancelled twice', async () => { + const id = await createSession(); + const tasks = await mainAgentTasks(id); + tasks.registerTask(fakeTask('agent', undefined, 'agent-ui')); + await flush(); + + await cancelTask(id, 'agent-ui'); + const repeated = await cancelTask(id, 'agent-ui'); + expect(repeated.body).toMatchObject({ + code: 40904, + data: { cancelled: false }, + details: { current_status: 'cancelled' }, + }); + }); + it('includes output_preview / output_bytes when with_output is set', async () => { const id = await createSession(); const tasks = await mainAgentTasks(id); @@ -253,17 +310,13 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { const taskId = tasks.registerTask(fakeTask('process')); await flush(); - const cancelled = await postJson<{ cancelled: boolean }>( - `/api/v1/sessions/${id}/tasks/${taskId}:cancel`, - ); + const cancelled = await cancelTask(id, taskId); expect(cancelled.body.code).toBe(0); expect(cancelled.body.data).toEqual({ cancelled: true }); // The task is now terminal (killed → cancelled); a second cancel is a // conflict with the idempotent envelope shape. - const again = await postJson<{ cancelled: boolean }>( - `/api/v1/sessions/${id}/tasks/${taskId}:cancel`, - ); + const again = await cancelTask(id, taskId); expect(again.body.code).toBe(40904); expect(again.body.data).toEqual({ cancelled: false }); expect(again.body.details).toEqual({ current_status: 'cancelled' }); @@ -272,10 +325,87 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { it('cancelling an unknown task returns 40406', async () => { const id = await createSession(); await mainAgentTasks(id); - const { body } = await postJson(`/api/v1/sessions/${id}/tasks/nope:cancel`); + const { body } = await cancelTask(id, 'nope'); expect(body.code).toBe(40406); }); + it('cancels a swarm member by Agent id and keeps repeated stop idempotent', async () => { + const id = await createSession(); + const swarm = await sessionSwarm(id); + const stopAgent = vi + .spyOn(swarm, 'stopAgent') + .mockReturnValueOnce({ + kind: 'stopping', + agentId: 'agent-swarm', + }) + .mockReturnValueOnce({ + kind: 'already_terminal', + agentId: 'agent-swarm', + status: 'aborted', + }); + + const cancelled = await cancelTask(id, 'agent-swarm'); + expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); + expect(stopAgent).toHaveBeenLastCalledWith({ + callerAgentId: 'main', + agentId: 'agent-swarm', + }); + + const again = await cancelTask(id, 'agent-swarm'); + expect(again.body).toMatchObject({ + code: 40904, + data: { cancelled: false }, + details: { current_status: 'cancelled' }, + }); + }); + + it('keeps resumed swarm cancellation authoritative over terminal task history', async () => { + const id = await createSession(); + const tasks = await mainAgentTasks(id); + const historicalTaskId = tasks.registerTask(completedAgentTask('agent-resumed')); + expect(await tasks.wait(historicalTaskId)).toMatchObject({ status: 'completed' }); + expect(tasks.getAgentTask('agent-resumed')).toMatchObject({ + taskId: historicalTaskId, + status: 'completed', + }); + + const swarm = await sessionSwarm(id); + const stopAgent = vi.spyOn(swarm, 'stopAgent').mockReturnValue({ + kind: 'stopping', + agentId: 'agent-resumed', + }); + + const cancelled = await cancelTask(id, 'agent-resumed'); + + expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); + expect(stopAgent).toHaveBeenCalledWith({ + callerAgentId: 'main', + agentId: 'agent-resumed', + }); + }); + + it('reports an exact terminal task id before a matching active swarm member', async () => { + const id = await createSession(); + const tasks = await mainAgentTasks(id); + const historicalTaskId = tasks.registerTask(completedAgentTask('agent-history')); + await tasks.wait(historicalTaskId); + + const swarm = await sessionSwarm(id); + const stopAgent = vi.spyOn(swarm, 'stopAgent').mockReturnValue({ + kind: 'stopping', + agentId: historicalTaskId, + }); + + const cancelled = await cancelTask(id, historicalTaskId); + + expect(cancelled.body).toMatchObject({ + code: 40904, + data: { cancelled: false }, + details: { current_status: 'completed' }, + }); + expect(stopAgent).not.toHaveBeenCalled(); + }); + it('rejects a bare POST without the :cancel suffix (40001)', async () => { const id = await createSession(); const tasks = await mainAgentTasks(id); @@ -293,7 +423,7 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { const got = await getJson('/api/v1/sessions/nope/tasks/tid'); expect(got.body.code).toBe(40401); - const cancelled = await postJson('/api/v1/sessions/nope/tasks/tid:cancel'); + const cancelled = await cancelTask('nope', 'tid'); expect(cancelled.body.code).toBe(40401); }); }); diff --git a/packages/protocol/src/__tests__/events.test.ts b/packages/protocol/src/__tests__/events.test.ts index 365fd7b9e6..38e9439c41 100644 --- a/packages/protocol/src/__tests__/events.test.ts +++ b/packages/protocol/src/__tests__/events.test.ts @@ -143,6 +143,23 @@ describe('events / display re-exports', () => { expect((parsed as { info: { detached?: boolean } }).info.detached).toBe(false); }); + it('preserves explicit subagent cancellation', () => { + const parsed = eventSchema.parse({ + type: 'subagent.failed', + agentId: 'main', + sessionId: 'sess_1', + subagentId: 'agent_1', + error: 'Aborted by the user', + cancelled: true, + }); + + expect(parsed).toMatchObject({ + type: 'subagent.failed', + subagentId: 'agent_1', + cancelled: true, + }); + }); + it('validates event.session.created events', () => { const parsed = eventSchema.parse({ type: 'event.session.created', diff --git a/packages/protocol/src/events.ts b/packages/protocol/src/events.ts index de09c1e65b..31822c302d 100644 --- a/packages/protocol/src/events.ts +++ b/packages/protocol/src/events.ts @@ -759,6 +759,8 @@ export interface SubagentFailedEvent { readonly type: 'subagent.failed'; readonly subagentId: string; readonly error: string; + /** True when the run ended through an explicit user cancellation. */ + readonly cancelled?: boolean; } export interface CompactionStartedEvent { @@ -1599,6 +1601,7 @@ export const subagentFailedEventSchema = z.object({ type: z.literal('subagent.failed'), subagentId: z.string(), error: z.string(), + cancelled: z.boolean().optional(), }) satisfies z.ZodType; export const compactionStartedEventSchema = z.object({ From 8daa22185d67d03ebbb6015666995256c3f0d7e8 Mon Sep 17 00:00:00 2001 From: qer Date: Tue, 14 Jul 2026 05:27:29 +0800 Subject: [PATCH 2/4] test: preserve swarm lifecycle mock signature --- .../agent-core-v2/test/session/swarm/sessionSwarm.test.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts b/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts index 2c2f6cee72..bb20a74f20 100644 --- a/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts +++ b/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts @@ -978,7 +978,7 @@ describe('SessionSwarmService metadata compatibility', () => { let agents: Record; let handles: Map; let lifecycle: IAgentLifecycleService; - let createAgent: ReturnType; + let createAgent: ReturnType>; let runAgent: ReturnType; let eventBus: IEventBus; @@ -989,7 +989,7 @@ describe('SessionSwarmService metadata compatibility', () => { handles = new Map(); eventBus = eventBusStub(); lifecycle = lifecycleStub(handles, eventBus); - createAgent = lifecycle.create as ReturnType; + createAgent = lifecycle.create as ReturnType>; runAgent = lifecycle.run as ReturnType; handles.set('main', agentHandle('main', lifecycle, eventBus)); @@ -1190,7 +1190,7 @@ describe('SessionSwarmService metadata compatibility', () => { [IAgentUserToolService, parentUserTools], ])), ); - createAgent.mockImplementationOnce((opts: CreateAgentOptions = {}) => { + createAgent.mockImplementationOnce(async (opts: CreateAgentOptions = {}) => { const id = opts.agentId ?? 'agent-new'; const handle = agentHandle( id, From 8678de1fe2a31194800ef96d537f001cf29965ae Mon Sep 17 00:00:00 2001 From: qer Date: Tue, 14 Jul 2026 05:46:16 +0800 Subject: [PATCH 3/4] fix: resolve swarm owner when stopping members --- .../src/session/swarm/sessionSwarm.ts | 5 +- .../src/session/swarm/sessionSwarmService.ts | 113 ++++++++---------- .../test/agent/swarm/swarm.test.ts | 4 +- packages/agent-core-v2/test/harness/agent.ts | 2 +- .../test/session/swarm/sessionSwarm.test.ts | 78 +++++++----- packages/agent-core-v2/test/tool/tool.test.ts | 10 +- packages/kap-server/src/routes/tasks.ts | 7 +- packages/kap-server/test/tasks.test.ts | 77 ++++++++++-- 8 files changed, 184 insertions(+), 112 deletions(-) diff --git a/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts b/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts index 1ad64f0aa1..1ae3ceeb65 100644 --- a/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts +++ b/packages/agent-core-v2/src/session/swarm/sessionSwarm.ts @@ -79,10 +79,7 @@ export interface ISessionSwarmService { readonly agentId: string; }): Promise; run(args: SessionSwarmRunArgs): Promise[]>; - stopAgent(args: { - readonly callerAgentId: string; - readonly agentId: string; - }): SessionSwarmStopResult; + stopAgent(agentId: string): SessionSwarmStopResult; cancel(args: { readonly callerAgentId: string }): void; } diff --git a/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts b/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts index d1d5f5987a..747363265d 100644 --- a/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts +++ b/packages/agent-core-v2/src/session/swarm/sessionSwarmService.ts @@ -4,8 +4,9 @@ * Runs a batch of agents on behalf of a caller agent: builds an * `AgentRunBatchLauncher` on top of the `agentLifecycle` primitives * (`create({ binding })`, `run`), drives the internal `AgentRunBatch` - * scheduler, and tracks each caller's live batches so cancellation can target - * one member or all work owned by the caller. The caller ↔ child association + * scheduler, and tracks each caller's live batches plus each active member's + * owning batch so cancellation can target one member or all work owned by the + * caller. The caller ↔ child association * is this domain's own business data: requester-side display facts * (`subagent.spawned` wire signals carrying the swarm's tool-call context, * `subagent.suspended` when a task is @@ -76,20 +77,17 @@ const RESUMED_PROFILE_FALLBACK = 'subagent'; const RECENT_SWARM_TERMINAL_LIMIT = 128; type InFlightBatch = { + readonly callerAgentId: string; readonly controller: AbortController; readonly batch: Pick, 'stopAgent'>; readonly agentIds: Set; }; -type CallerInFlight = { - readonly batches: Set; - readonly byAgentId: Map; -}; - export class SessionSwarmService implements ISessionSwarmService { declare readonly _serviceBrand: undefined; - private readonly inFlightByCaller = new Map(); + private readonly inFlightByCaller = new Map>(); + private readonly inFlightByAgentId = new Map(); private readonly recentTerminalStatuses = new Map< string, SessionSwarmRunResult['status'] @@ -140,18 +138,18 @@ export class SessionSwarmService implements ISessionSwarmService { }; const maxConcurrency = resolveSwarmMaxConcurrency(); const batch = new AgentRunBatch(launcher, linkedTasks, { maxConcurrency }); - inFlight = { controller, batch, agentIds: new Set() }; - this.addInFlight(callerAgentId, inFlight); + inFlight = { callerAgentId, controller, batch, agentIds: new Set() }; + this.addInFlight(inFlight); const promise = batch.run(); const cleanup = () => { for (const unlink of unlinks) unlink(); - this.removeInFlight(callerAgentId, inFlight); + this.removeInFlight(inFlight); }; void promise.then( (results) => { for (const result of results) { if (result.agentId === undefined) continue; - this.rememberTerminal(callerAgentId, result.agentId, result.status); + this.rememberTerminal(result.agentId, result.status); } cleanup(); }, @@ -162,40 +160,36 @@ export class SessionSwarmService implements ISessionSwarmService { return promise; } - stopAgent(args: { - readonly callerAgentId: string; - readonly agentId: string; - }): SessionSwarmStopResult { - const caller = this.inFlightByCaller.get(args.callerAgentId); - const inFlight = caller?.byAgentId.get(args.agentId); - const current = inFlight?.batch.stopAgent(args.agentId); - if (current !== undefined && current.kind !== 'not_found') { - if (current.kind === 'stopped') { - const callerHandle = this.lifecycle.getHandle(args.callerAgentId); - if (callerHandle !== undefined) { - this.publishCancellation(callerHandle, args.agentId); + stopAgent(agentId: string): SessionSwarmStopResult { + const inFlight = this.inFlightByAgentId.get(agentId); + if (inFlight !== undefined) { + const current = inFlight.batch.stopAgent(agentId); + if (current.kind !== 'not_found') { + if (current.kind === 'stopped') { + const callerHandle = this.lifecycle.getHandle(inFlight.callerAgentId); + if (callerHandle !== undefined) { + this.publishCancellation(callerHandle, agentId); + } } + return current; } - return current; } - const status = this.recentTerminalStatus( - this.agentRunKey(args.callerAgentId, args.agentId), - ); + const status = this.recentTerminalStatus(agentId); if (status !== undefined) { return { kind: 'already_terminal', - agentId: args.agentId, + agentId, status, }; } - return { kind: 'not_found', agentId: args.agentId }; + return { kind: 'not_found', agentId }; } cancel({ callerAgentId }: { readonly callerAgentId: string }): void { - const caller = this.inFlightByCaller.get(callerAgentId); - if (caller === undefined) return; - for (const inFlight of caller.batches) { + const batches = this.inFlightByCaller.get(callerAgentId); + if (batches === undefined) return; + for (const inFlight of batches) { inFlight.controller.abort(userCancellationReason()); } } @@ -343,17 +337,17 @@ export class SessionSwarmService implements ISessionSwarmService { agentId: string, options: AgentRunAttemptOptions, ): void { - const caller = this.inFlightByCaller.get(callerAgentId); - if (caller === undefined || !caller.batches.has(inFlight)) { + const batches = this.inFlightByCaller.get(callerAgentId); + if (batches === undefined || !batches.has(inFlight)) { options.signal.throwIfAborted(); throw new Error('Swarm batch is no longer running'); } - if (caller.byAgentId.has(agentId)) { + if (this.inFlightByAgentId.has(agentId)) { throw new Error(`Agent instance "${agentId}" is already owned by a running swarm batch`); } - this.recentTerminalStatuses.delete(this.agentRunKey(callerAgentId, agentId)); + this.recentTerminalStatuses.delete(agentId); inFlight.agentIds.add(agentId); - caller.byAgentId.set(agentId, inFlight); + this.inFlightByAgentId.set(agentId, inFlight); options.onAgentIdentified?.(agentId); } @@ -375,39 +369,42 @@ export class SessionSwarmService implements ISessionSwarmService { agentId: string, signal: AbortSignal, ): void { - const caller = this.inFlightByCaller.get(callerAgentId); - if (caller?.byAgentId.get(agentId) === inFlight) return; + if ( + inFlight.callerAgentId === callerAgentId && + this.inFlightByAgentId.get(agentId) === inFlight + ) { + return; + } signal.throwIfAborted(); throw new Error(`Agent instance "${agentId}" is not owned by this swarm batch`); } - private addInFlight(callerAgentId: string, inFlight: InFlightBatch): void { - let caller = this.inFlightByCaller.get(callerAgentId); - if (caller === undefined) { - caller = { batches: new Set(), byAgentId: new Map() }; - this.inFlightByCaller.set(callerAgentId, caller); + private addInFlight(inFlight: InFlightBatch): void { + let batches = this.inFlightByCaller.get(inFlight.callerAgentId); + if (batches === undefined) { + batches = new Set(); + this.inFlightByCaller.set(inFlight.callerAgentId, batches); } - caller.batches.add(inFlight); + batches.add(inFlight); } - private removeInFlight(callerAgentId: string, inFlight: InFlightBatch): void { - const caller = this.inFlightByCaller.get(callerAgentId); - if (caller === undefined) return; - caller.batches.delete(inFlight); + private removeInFlight(inFlight: InFlightBatch): void { + const batches = this.inFlightByCaller.get(inFlight.callerAgentId); + batches?.delete(inFlight); for (const agentId of inFlight.agentIds) { - if (caller.byAgentId.get(agentId) === inFlight) caller.byAgentId.delete(agentId); + if (this.inFlightByAgentId.get(agentId) === inFlight) { + this.inFlightByAgentId.delete(agentId); + } } - if (caller.batches.size === 0) this.inFlightByCaller.delete(callerAgentId); + if (batches?.size === 0) this.inFlightByCaller.delete(inFlight.callerAgentId); } private rememberTerminal( - callerAgentId: string, agentId: string, status: SessionSwarmRunResult['status'], ): void { - const key = this.agentRunKey(callerAgentId, agentId); - this.recentTerminalStatuses.delete(key); - this.recentTerminalStatuses.set(key, status); + this.recentTerminalStatuses.delete(agentId); + this.recentTerminalStatuses.set(agentId, status); while (this.recentTerminalStatuses.size > RECENT_SWARM_TERMINAL_LIMIT) { const oldest = this.recentTerminalStatuses.keys().next().value; if (oldest === undefined) break; @@ -425,10 +422,6 @@ export class SessionSwarmService implements ISessionSwarmService { return status; } - private agentRunKey(callerAgentId: string, agentId: string): string { - return `${callerAgentId}\0${agentId}`; - } - private requireHandle(agentId: string, label: string): IAgentScopeHandle { const handle = this.lifecycle.getHandle(agentId); if (handle === undefined) throw new Error(`${label} "${agentId}" does not exist`); diff --git a/packages/agent-core-v2/test/agent/swarm/swarm.test.ts b/packages/agent-core-v2/test/agent/swarm/swarm.test.ts index bca2d1ea37..096d1a101a 100644 --- a/packages/agent-core-v2/test/agent/swarm/swarm.test.ts +++ b/packages/agent-core-v2/test/agent/swarm/swarm.test.ts @@ -57,7 +57,7 @@ function mockSwarmHost({ _serviceBrand: undefined, getSwarmItem, run, - stopAgent: ({ agentId }: { agentId: string }) => ({ kind: 'not_found' as const, agentId }), + stopAgent: (agentId: string) => ({ kind: 'not_found' as const, agentId }), cancel: vi.fn(), }, callerAgentId: 'main', @@ -90,7 +90,7 @@ describe('AgentSwarmService', () => { ix.stub(ISessionSwarmService, { getSwarmItem: async () => undefined, run: async () => [], - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, }); ix.stub(IAgentScopeContext, makeAgentScopeContext({ agentId: 'main', agentScope: '' })); diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index fab49f0886..c35f1377fa 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -654,7 +654,7 @@ export function swarmServices( _serviceBrand: undefined, getSwarmItem: async () => undefined, run: swarmService, - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, } satisfies ISessionSwarmService : swarmService; diff --git a/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts b/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts index bb20a74f20..f8f1b75387 100644 --- a/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts +++ b/packages/agent-core-v2/test/session/swarm/sessionSwarm.test.ts @@ -1070,9 +1070,13 @@ describe('SessionSwarmService metadata compatibility', () => { throw new Error(`Unexpected swarm run for ${agentId}`); } signals.set(agentId, signal); - signal.addEventListener('abort', () => completion.reject(signal.reason), { - once: true, - }); + signal.addEventListener( + 'abort', + () => { + completion.reject(signal.reason); + }, + { once: true }, + ); options.onReady?.(); if (signals.size === agentIds.length) started.resolve(); return { agentId, turn: {} as never, completion }; @@ -1361,7 +1365,7 @@ describe('SessionSwarmService metadata compatibility', () => { await vi.advanceTimersByTimeAsync(0); expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-rate-limited' }), + service.stopAgent('agent-rate-limited'), ).toMatchObject({ kind: 'stopped' }); expect(published).toContainEqual( expect.objectContaining({ @@ -1388,7 +1392,9 @@ describe('SessionSwarmService metadata compatibility', () => { options?.onReady?.(); options?.signal.addEventListener( 'abort', - () => completion.reject(options.signal.reason), + () => { + completion.reject(options.signal.reason); + }, { once: true }, ); return { agentId, turn: {} as never, completion }; @@ -1398,9 +1404,11 @@ describe('SessionSwarmService metadata compatibility', () => { callerAgentId: 'main', tasks: [spawnSessionTask('src/a.ts')], }); - await vi.waitFor(() => expect(runAgent).toHaveBeenCalledOnce()); + await vi.waitFor(() => { + expect(runAgent).toHaveBeenCalledOnce(); + }); - expect(service.stopAgent({ callerAgentId: 'main', agentId: 'agent-new' })).toEqual({ + expect(service.stopAgent('agent-new')).toEqual({ kind: 'stopping', agentId: 'agent-new', }); @@ -1408,7 +1416,7 @@ describe('SessionSwarmService metadata compatibility', () => { { agentId: 'agent-new', status: 'aborted', state: 'started' }, ]); - expect(service.stopAgent({ callerAgentId: 'main', agentId: 'agent-new' })).toEqual({ + expect(service.stopAgent('agent-new')).toEqual({ kind: 'already_terminal', agentId: 'agent-new', status: 'aborted', @@ -1430,7 +1438,7 @@ describe('SessionSwarmService metadata compatibility', () => { await runs.started; expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-first-batch' }), + service.stopAgent('agent-first-batch'), ).toEqual({ kind: 'stopping', agentId: 'agent-first-batch', @@ -1474,7 +1482,7 @@ describe('SessionSwarmService metadata compatibility', () => { }, ]); expect(runAgent).toHaveBeenCalledOnce(); - expect(service.stopAgent({ callerAgentId: 'main', agentId: 'agent-shared' })).toEqual({ + expect(service.stopAgent('agent-shared')).toEqual({ kind: 'stopping', agentId: 'agent-shared', }); @@ -1553,10 +1561,10 @@ describe('SessionSwarmService metadata compatibility', () => { } expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-terminal-0' }), + service.stopAgent('agent-terminal-0'), ).toEqual({ kind: 'not_found', agentId: 'agent-terminal-0' }); expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-terminal-128' }), + service.stopAgent('agent-terminal-128'), ).toEqual({ kind: 'already_terminal', agentId: 'agent-terminal-128', @@ -1576,7 +1584,7 @@ describe('SessionSwarmService metadata compatibility', () => { }); } expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-lru-0' }), + service.stopAgent('agent-lru-0'), ).toEqual({ kind: 'already_terminal', agentId: 'agent-lru-0', @@ -1589,10 +1597,10 @@ describe('SessionSwarmService metadata compatibility', () => { }); expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-lru-1' }), + service.stopAgent('agent-lru-1'), ).toEqual({ kind: 'not_found', agentId: 'agent-lru-1' }); expect( - service.stopAgent({ callerAgentId: 'main', agentId: 'agent-lru-0' }), + service.stopAgent('agent-lru-0'), ).toEqual({ kind: 'already_terminal', agentId: 'agent-lru-0', @@ -1600,31 +1608,47 @@ describe('SessionSwarmService metadata compatibility', () => { }); }); - it('does not stop an in-flight member owned by another caller', async () => { + it('stops an in-flight member through its recorded caller ownership', async () => { const completion = createControlledPromise<{ summary: string }>(); let runSignal: AbortSignal | undefined; + const otherEventBus = eventBusStub(); + handles.set('other', agentHandle('other', lifecycle, otherEventBus)); runAgent.mockImplementation((agentId, _request, options) => { runSignal = options?.signal; + options?.signal.addEventListener( + 'abort', + () => { + completion.reject(options.signal.reason); + }, + { once: true }, + ); options?.onReady?.(); return { agentId, turn: {} as never, completion }; }); const service = ix.get(ISessionSwarmService); const running = service.run({ - callerAgentId: 'main', + callerAgentId: 'other', tasks: [spawnSessionTask('src/a.ts')], }); - await vi.waitFor(() => expect(runAgent).toHaveBeenCalledOnce()); + await vi.waitFor(() => { + expect(runAgent).toHaveBeenCalledOnce(); + }); - expect(service.stopAgent({ callerAgentId: 'other', agentId: 'agent-new' })).toEqual({ - kind: 'not_found', + expect(service.stopAgent('agent-new')).toEqual({ + kind: 'stopping', agentId: 'agent-new', }); - expect(runSignal?.aborted).toBe(false); - - completion.resolve({ summary: 'child summary' }); + expect(runSignal?.aborted).toBe(true); await expect(running).resolves.toMatchObject([ - { agentId: 'agent-new', status: 'completed' }, + { agentId: 'agent-new', status: 'aborted' }, ]); + expect(otherEventBus.publish).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'subagent.failed', + subagentId: 'agent-new', + cancelled: true, + }), + ); }); it('rejects resume of an already running child before launching or emitting spawned', async () => { @@ -1844,7 +1868,7 @@ function createMockAgentRunBatchRunner( readonly stopAgent: (agentId: string) => SessionSwarmStopResult; } { const attempts: MockAgentRunAttemptRecord[] = []; - let activeTasks: readonly QueuedAgentRunTask[] = []; + let activeTasks: readonly QueuedAgentRunTask[] = []; let activeBatch: AgentRunBatch | undefined; const createHandle = ( @@ -1914,7 +1938,7 @@ function createMockAgentRunBatchRunner( } function findMockAgentRunTask( - tasks: readonly QueuedAgentRunTask[], + tasks: readonly QueuedAgentRunTask[], options: AgentRunAttemptOptions, ): QueuedAgentRunTask { const task = tasks.find( @@ -1928,7 +1952,7 @@ function findMockAgentRunTask( return task as QueuedAgentRunTask; } -function mockAgentRunId(task: QueuedAgentRunTask, attemptIndex: number): string { +function mockAgentRunId(task: QueuedAgentRunTask, attemptIndex: number): string { if (typeof task.data === 'number') return `agent-${String(task.data)}`; return `agent-${String(attemptIndex + 1)}`; } diff --git a/packages/agent-core-v2/test/tool/tool.test.ts b/packages/agent-core-v2/test/tool/tool.test.ts index bf29006225..6b1a51dd4a 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -1453,7 +1453,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1535,7 +1535,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem, run: runSwarm as ISessionSwarmService['run'], - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1661,7 +1661,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1709,7 +1709,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); @@ -1766,7 +1766,7 @@ describe('AgentSwarm tool execution contract', () => { _serviceBrand: undefined, getSwarmItem: async () => undefined, run: runSwarm as ISessionSwarmService['run'], - stopAgent: ({ agentId }) => ({ kind: 'not_found', agentId }), + stopAgent: (agentId) => ({ kind: 'not_found', agentId }), cancel: () => {}, }; ctx = createTestAgent(swarmServices(swarmService)); diff --git a/packages/kap-server/src/routes/tasks.ts b/packages/kap-server/src/routes/tasks.ts index 4f9f7b354e..c303655cfd 100644 --- a/packages/kap-server/src/routes/tasks.ts +++ b/packages/kap-server/src/routes/tasks.ts @@ -63,7 +63,7 @@ import { z } from 'zod'; import { errEnvelope, okEnvelope } from '../envelope'; import { defineRoute } from '../middleware/defineRoute'; -import { ensureMainAgent, MAIN_AGENT_ID } from '../transport/mainAgent'; +import { ensureMainAgent } from '../transport/mainAgent'; import { parseActionSuffix } from './action-suffix'; /** Default cap (bytes) for the opt-in output preview on GET-by-id. */ @@ -332,10 +332,7 @@ function resolveCancelTarget( return { kind: 'agent_task', task: agentTask }; } - const swarmResult = resolved.swarm?.stopAgent({ - callerAgentId: MAIN_AGENT_ID, - agentId: taskOrAgentId, - }); + const swarmResult = resolved.swarm?.stopAgent(taskOrAgentId); if (swarmResult !== undefined && swarmResult.kind !== 'not_found') { return { kind: 'swarm', result: swarmResult }; } diff --git a/packages/kap-server/test/tasks.test.ts b/packages/kap-server/test/tasks.test.ts index 4487b8366c..7fffba3e01 100644 --- a/packages/kap-server/test/tasks.test.ts +++ b/packages/kap-server/test/tasks.test.ts @@ -4,6 +4,7 @@ import { join } from 'node:path'; import { IAgentLifecycleService, + IAgentProfileService, IAgentTaskService, ISessionLifecycleService, ISessionSwarmService, @@ -346,10 +347,7 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { const cancelled = await cancelTask(id, 'agent-swarm'); expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); - expect(stopAgent).toHaveBeenLastCalledWith({ - callerAgentId: 'main', - agentId: 'agent-swarm', - }); + expect(stopAgent).toHaveBeenLastCalledWith('agent-swarm'); const again = await cancelTask(id, 'agent-swarm'); expect(again.body).toMatchObject({ @@ -359,6 +357,72 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { }); }); + it('cancels through REST when a non-main agent owns the swarm member', async () => { + const id = await createSession(); + const swarm = await sessionSwarm(id); + const session = server!.core.accessor.get(ISessionLifecycleService).get(id); + if (session === undefined) throw new Error(`session ${id} not found`); + const lifecycle = session.accessor.get(IAgentLifecycleService); + const caller = await lifecycle.create({ agentId: 'agent-caller' }); + caller.accessor.get(IAgentProfileService).update({ modelAlias: 'test-model' }); + await lifecycle.create({ + agentId: 'agent-owned-child', + labels: { parentAgentId: caller.id }, + }); + + let rejectCompletion: (reason?: unknown) => void = () => {}; + const completion = new Promise<{ summary: string }>((_resolve, reject) => { + rejectCompletion = reject; + }); + let runSignal: AbortSignal | undefined; + const runAgent = vi.spyOn(lifecycle, 'run').mockImplementation( + async (agentId, _request, options) => { + runSignal = options.signal; + options.signal.addEventListener( + 'abort', + () => { + rejectCompletion(options.signal.reason); + }, + { once: true }, + ); + options.onReady?.(); + return { agentId, turn: {} as never, completion }; + }, + ); + const running = swarm.run({ + callerAgentId: caller.id, + tasks: [ + { + kind: 'resume', + data: undefined, + profileName: 'subagent', + parentToolCallId: 'call_non_main_swarm', + prompt: 'Continue', + description: 'Continue child', + runInBackground: false, + resumeAgentId: 'agent-owned-child', + }, + ], + }); + await vi.waitFor(() => { + expect(runAgent).toHaveBeenCalledOnce(); + }); + + const cancelled = await cancelTask(id, 'agent-owned-child'); + + expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); + expect(runSignal?.aborted).toBe(true); + await expect(running).resolves.toMatchObject([ + { agentId: 'agent-owned-child', status: 'aborted' }, + ]); + const repeated = await cancelTask(id, 'agent-owned-child'); + expect(repeated.body).toMatchObject({ + code: 40904, + data: { cancelled: false }, + details: { current_status: 'cancelled' }, + }); + }); + it('keeps resumed swarm cancellation authoritative over terminal task history', async () => { const id = await createSession(); const tasks = await mainAgentTasks(id); @@ -378,10 +442,7 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { const cancelled = await cancelTask(id, 'agent-resumed'); expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); - expect(stopAgent).toHaveBeenCalledWith({ - callerAgentId: 'main', - agentId: 'agent-resumed', - }); + expect(stopAgent).toHaveBeenCalledWith('agent-resumed'); }); it('reports an exact terminal task id before a matching active swarm member', async () => { From 1b7551dd4296f88740dc4e3733c7e8ea37367f2d Mon Sep 17 00:00:00 2001 From: qer Date: Tue, 14 Jul 2026 06:37:30 +0800 Subject: [PATCH 4/4] fix(agent-core-v2): publish start-hook cancellations --- .../session/agentLifecycle/mirrorAgentRun.ts | 35 +++-- packages/agent-core-v2/test/tool/tool.test.ts | 122 +++++++++++++++--- packages/kap-server/test/tasks.test.ts | 39 +++++- 3 files changed, 165 insertions(+), 31 deletions(-) diff --git a/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts b/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts index 49f71d37f4..8043e436d1 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/mirrorAgentRun.ts @@ -27,6 +27,7 @@ import type { IAgentScopeHandle } from '#/_base/di/scope'; import { isAbortError, isUserCancellation, + type UserCancellationError, userCancellationReason, } from '#/_base/utils/abort'; import { IAgentContextSizeService } from '#/agent/contextSize/contextSize'; @@ -118,11 +119,30 @@ export async function mirrorAgentRun( ): Promise<{ summary: string; usage?: TokenUsage }> { const eventBus = requester.accessor.get(IEventBus); const agentLifecycle = requester.accessor.get(IAgentLifecycleService); + let cancellationPublished = false; + const resolveCancellation = (reason: unknown): UserCancellationError | undefined => + isUserCancellation(reason) + ? reason + : isUserCancellation(options.signal.reason) + ? options.signal.reason + : undefined; + const publishCancellation = (reason: UserCancellationError): void => { + if (cancellationPublished) return; + cancellationPublished = true; + eventBus?.publish({ + type: 'subagent.failed', + subagentId: run.agentId, + error: errorMessage(reason), + cancelled: true, + }); + }; + void run.completion.catch(() => {}); eventBus?.publish({ type: 'subagent.started', subagentId: run.agentId }); if (options.prompt !== undefined) { const cancelAndRethrow = (reason: unknown): never => { + const cancellationReason = resolveCancellation(reason); + if (cancellationReason !== undefined) publishCancellation(cancellationReason); options.cancel?.(reason); - void run.completion.catch(() => {}); throw reason; }; try { @@ -154,18 +174,9 @@ export async function mirrorAgentRun( }); return result; } catch (error) { - const cancellationReason = isUserCancellation(error) - ? error - : isUserCancellation(options.signal.reason) - ? options.signal.reason - : undefined; + const cancellationReason = resolveCancellation(error); if (cancellationReason !== undefined) { - eventBus?.publish({ - type: 'subagent.failed', - subagentId: run.agentId, - error: errorMessage(cancellationReason), - cancelled: true, - }); + publishCancellation(cancellationReason); } else if (!isAbortError(error) && !shouldSuppressFailure(options, error)) { eventBus?.publish({ type: 'subagent.failed', diff --git a/packages/agent-core-v2/test/tool/tool.test.ts b/packages/agent-core-v2/test/tool/tool.test.ts index 6b1a51dd4a..2c769b2484 100644 --- a/packages/agent-core-v2/test/tool/tool.test.ts +++ b/packages/agent-core-v2/test/tool/tool.test.ts @@ -287,6 +287,41 @@ function createAgentLifecycleStub(options: AgentLifecycleStubOptions = {}): Agen return lifecycle; } +function createMirrorRequester(lifecycle: IAgentLifecycleService): { + readonly requester: IAgentScopeHandle; + readonly events: DomainEvent[]; +} { + const events: DomainEvent[] = []; + const eventBus = { + _serviceBrand: undefined, + publish: vi.fn((event: DomainEvent) => { + events.push(event); + }), + subscribe: vi.fn(() => noopDisposable()), + } as IEventBus; + return { + events, + requester: { + id: 'main', + kind: LifecycleScope.Agent, + accessor: { + get: ((serviceId: unknown) => { + if (serviceId === IEventBus) return eventBus; + if (serviceId === IAgentLifecycleService) return lifecycle; + return undefined; + }) as IAgentScopeHandle['accessor']['get'], + }, + dispose: () => {}, + }, + }; +} + +function cancelledFailures(events: readonly DomainEvent[]): DomainEvent[] { + return events.filter( + (event) => event.type === 'subagent.failed' && event.cancelled === true, + ); +} + function agentTool(ctx: TestAgentContext): ExecutableTool { const tool = ctx.get(IAgentToolRegistryService).resolve('Agent'); expect(tool).toBeDefined(); @@ -652,24 +687,7 @@ describe('Agent tool execution contract', () => { it('mirrors explicit user cancellation as a cancelled subagent failure', async () => { const lifecycle = createAgentLifecycleStub(); - const events: DomainEvent[] = []; - const eventBus = { - _serviceBrand: undefined, - publish: vi.fn((event: DomainEvent) => events.push(event)), - subscribe: vi.fn(() => noopDisposable()), - } as IEventBus; - const requester = { - id: 'main', - kind: LifecycleScope.Agent, - accessor: { - get: ((serviceId: unknown) => { - if (serviceId === IEventBus) return eventBus; - if (serviceId === IAgentLifecycleService) return lifecycle; - return undefined; - }) as IAgentScopeHandle['accessor']['get'], - }, - dispose: () => {}, - } satisfies IAgentScopeHandle; + const { requester, events } = createMirrorRequester(lifecycle); const controller = new AbortController(); const reason = userCancellationReason(); controller.abort(reason); @@ -694,6 +712,74 @@ describe('Agent tool execution contract', () => { }); }); + it('publishes one cancelled failure when the signal is aborted before the start hook', async () => { + const lifecycle = createAgentLifecycleStub(); + const { requester, events } = createMirrorRequester(lifecycle); + const controller = new AbortController(); + const reason = userCancellationReason(); + const completion = deferred<{ summary: string }>(); + const cancel = vi.fn((cancelReason?: unknown) => { + completion.reject(cancelReason); + }); + controller.abort(reason); + + const mirrored = mirrorAgentRun( + requester, + { + agentId: 'agent-child', + turn: {} as AgentRunHandle['turn'], + completion: completion.promise, + }, + { + profileName: 'explore', + prompt: 'Investigate', + signal: controller.signal, + cancel, + }, + ); + + await expect(mirrored).rejects.toBe(reason); + expect(lifecycle.hooks.onWillStartAgentTask.run).toHaveBeenCalledOnce(); + expect(cancel).toHaveBeenCalledWith(reason); + expect(cancelledFailures(events)).toEqual([ + expect.objectContaining({ subagentId: 'agent-child', cancelled: true }), + ]); + }); + + it('publishes one cancelled failure when the signal is aborted after the start hook', async () => { + const lifecycle = createAgentLifecycleStub(); + const { requester, events } = createMirrorRequester(lifecycle); + const hookFinished = deferred(); + vi.mocked(lifecycle.hooks.onWillStartAgentTask.run).mockImplementation(async () => { + hookFinished.resolve(undefined); + }); + const controller = new AbortController(); + const completion = deferred<{ summary: string }>(); + const mirrored = mirrorAgentRun( + requester, + { + agentId: 'agent-child', + turn: {} as AgentRunHandle['turn'], + completion: completion.promise, + }, + { + profileName: 'explore', + prompt: 'Investigate', + signal: controller.signal, + }, + ); + await hookFinished.promise; + await Promise.resolve(); + const reason = userCancellationReason(); + controller.abort(reason); + completion.reject(reason); + + await expect(mirrored).rejects.toBe(reason); + expect(cancelledFailures(events)).toEqual([ + expect.objectContaining({ subagentId: 'agent-child', cancelled: true }), + ]); + }); + it('inherits parent user tools when spawning a subagent', async () => { const lookupTool: UserToolRegistration = { name: 'Lookup', diff --git a/packages/kap-server/test/tasks.test.ts b/packages/kap-server/test/tasks.test.ts index 7fffba3e01..0d30b7f6d9 100644 --- a/packages/kap-server/test/tasks.test.ts +++ b/packages/kap-server/test/tasks.test.ts @@ -3,6 +3,7 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { + IEventBus, IAgentLifecycleService, IAgentProfileService, IAgentTaskService, @@ -10,6 +11,7 @@ import { ISessionSwarmService, IModelResolver, type AgentTask, + type DomainEvent, } from '@moonshot-ai/agent-core-v2'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; @@ -357,7 +359,7 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { }); }); - it('cancels through REST when a non-main agent owns the swarm member', async () => { + it('publishes one cancellation when REST stops a non-main swarm during its start hook', async () => { const id = await createSession(); const swarm = await sessionSwarm(id); const session = server!.core.accessor.get(ISessionLifecycleService).get(id); @@ -370,6 +372,29 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { labels: { parentAgentId: caller.id }, }); + let enterHook: () => void = () => {}; + const hookEntered = new Promise((resolve) => { + enterHook = resolve; + }); + let releaseHook: () => void = () => {}; + const hookReleased = new Promise((resolve) => { + releaseHook = resolve; + }); + const hookRegistration = lifecycle.hooks.onWillStartAgentTask.register( + 'tasks-test-deferred-start', + async (_context, next) => { + enterHook(); + await hookReleased; + await next(); + }, + ); + const failures: Array> = []; + const eventRegistration = caller.accessor + .get(IEventBus) + .subscribe('subagent.failed', (event) => { + if (event.subagentId === 'agent-owned-child') failures.push(event); + }); + let rejectCompletion: (reason?: unknown) => void = () => {}; const completion = new Promise<{ summary: string }>((_resolve, reject) => { rejectCompletion = reject; @@ -407,20 +432,32 @@ describe('server-v2 /api/v1/sessions/{sid}/tasks', () => { await vi.waitFor(() => { expect(runAgent).toHaveBeenCalledOnce(); }); + await hookEntered; const cancelled = await cancelTask(id, 'agent-owned-child'); expect(cancelled.body).toMatchObject({ code: 0, data: { cancelled: true } }); expect(runSignal?.aborted).toBe(true); + expect(failures).toEqual([]); + releaseHook(); await expect(running).resolves.toMatchObject([ { agentId: 'agent-owned-child', status: 'aborted' }, ]); + expect(failures).toEqual([ + expect.objectContaining({ + subagentId: 'agent-owned-child', + cancelled: true, + }), + ]); const repeated = await cancelTask(id, 'agent-owned-child'); expect(repeated.body).toMatchObject({ code: 40904, data: { cancelled: false }, details: { current_status: 'cancelled' }, }); + expect(failures).toHaveLength(1); + eventRegistration.dispose(); + hookRegistration.dispose(); }); it('keeps resumed swarm cancellation authoritative over terminal task history', async () => {