From 45113d11df35f4c637be3f240e41aa4da78888ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Sun, 27 Sep 2026 19:09:30 +0800 Subject: [PATCH] fix: finalize stopped teammates when roster removal fails --- src/utils/swarm/spawnInProcess.test.ts | 99 ++++++++++++++++++++++++++ src/utils/swarm/spawnInProcess.ts | 52 +++++++------- 2 files changed, 127 insertions(+), 24 deletions(-) create mode 100644 src/utils/swarm/spawnInProcess.test.ts diff --git a/src/utils/swarm/spawnInProcess.test.ts b/src/utils/swarm/spawnInProcess.test.ts new file mode 100644 index 00000000..4d67098a --- /dev/null +++ b/src/utils/swarm/spawnInProcess.test.ts @@ -0,0 +1,99 @@ +import { expect, spyOn, test } from 'bun:test' +import { mkdtemp, rm } from 'fs/promises' +import { tmpdir } from 'os' +import { join } from 'path' +import type { AppState } from '../../state/AppState.js' +import * as sdk from '../sdkEventQueue.js' +import * as diskOutput from '../task/diskOutput.js' +import * as framework from '../task/framework.js' +import * as tracing from '../telemetry/perfettoTracing.js' +import { killInProcessTeammate } from './spawnInProcess.js' +import * as teamHelpers from './teamHelpers.js' + +async function withKillFixture(run: (fixture: { + setAppState: (update: (state: AppState) => AppState) => void + getState: () => AppState + abortController: AbortController + remove: ReturnType> + terminated: ReturnType> + evictOutput: ReturnType> + unregister: ReturnType> + timer: ReturnType> +}) => Promise) { + const directory = await mkdtemp(join(tmpdir(), 'cc-haha-kill-teammate-')) + const originalHome = process.env.HOME + const originalConfig = process.env.CLAUDE_CONFIG_DIR + process.env.HOME = directory + process.env.CLAUDE_CONFIG_DIR = join(directory, 'claude') + const abortController = new AbortController() + let state = { + tasks: { + worker: { + id: 'worker', + type: 'in_process_teammate', + status: 'running', + identity: { teamName: 'fixture', agentId: 'worker@fixture' }, + description: 'fixture teammate', + toolUseId: 'fixture-tool', + abortController, + }, + }, + teamContext: { teammates: { 'worker@fixture': {} } }, + } as unknown as AppState + const remove = spyOn(teamHelpers, 'removeMemberByAgentId').mockResolvedValue(true) + const terminated = spyOn(sdk, 'emitTaskTerminatedSdk').mockImplementation(() => {}) + const evictOutput = spyOn(diskOutput, 'evictTaskOutput').mockResolvedValue(undefined) + const unregister = spyOn(tracing, 'unregisterAgent').mockImplementation(() => {}) + const timer = spyOn(globalThis, 'setTimeout').mockImplementation( + (() => 0 as unknown as ReturnType) as typeof setTimeout, + ) + try { + await run({ + setAppState: update => { state = update(state) }, + getState: () => state, + abortController, + remove, + terminated, + evictOutput, + unregister, + timer, + }) + } finally { + for (const mock of [remove, terminated, evictOutput, unregister, timer]) mock.mockRestore() + if (originalHome === undefined) delete process.env.HOME + else process.env.HOME = originalHome + if (originalConfig === undefined) delete process.env.CLAUDE_CONFIG_DIR + else process.env.CLAUDE_CONFIG_DIR = originalConfig + await rm(directory, { recursive: true, force: true }) + } +} + +for (const failureCode of [undefined, 'ELOCKED', 'EPERM']) { + test(`killing a teammate finalizes exactly once when removal ${failureCode ?? 'succeeds'}`, async () => { + await withKillFixture(async fixture => { + const failure = Object.assign(new Error('fixture removal failure'), { code: failureCode }) + if (failureCode) fixture.remove.mockRejectedValue(failure) + const kill = killInProcessTeammate('worker', fixture.setAppState) + if (failureCode) await expect(kill).rejects.toBe(failure) + else expect(await kill).toBe(true) + + expect(fixture.abortController.signal.aborted).toBe(true) + expect(fixture.getState().tasks.worker).toMatchObject({ status: 'killed', notified: true }) + expect(fixture.getState().teamContext?.teammates).toEqual({}) + expect(fixture.remove).toHaveBeenCalledWith('fixture', 'worker@fixture') + expect(fixture.terminated).toHaveBeenCalledWith('worker', 'stopped', { + toolUseId: 'fixture-tool', summary: 'fixture teammate', ownerAgentId: 'worker@fixture', + }) + expect(fixture.evictOutput).toHaveBeenCalledWith('worker') + expect(fixture.unregister).toHaveBeenCalledWith('worker@fixture') + expect(fixture.timer).toHaveBeenCalledWith(expect.any(Function), framework.STOPPED_DISPLAY_MS) + + // A repeated stop must neither retry persistence nor duplicate finalizers. + expect(await killInProcessTeammate('worker', fixture.setAppState)).toBe(false) + expect(await killInProcessTeammate('missing', fixture.setAppState)).toBe(false) + for (const mock of [fixture.remove, fixture.terminated, fixture.evictOutput, fixture.unregister, fixture.timer]) { + expect(mock).toHaveBeenCalledTimes(1) + } + }) + }) +} diff --git a/src/utils/swarm/spawnInProcess.ts b/src/utils/swarm/spawnInProcess.ts index e785a5eb..0966311c 100644 --- a/src/utils/swarm/spawnInProcess.ts +++ b/src/utils/swarm/spawnInProcess.ts @@ -303,31 +303,35 @@ export async function killInProcessTeammate( } }) - // Remove from team file (outside state updater to avoid file I/O in callback) - if (teamName && agentId) { - await removeMemberByAgentId(teamName, agentId) - } + try { + // Remove from team file (outside state updater to avoid file I/O in callback) + if (teamName && agentId) { + await removeMemberByAgentId(teamName, agentId) + } + } finally { + // The task is already stopped. Persistence failures must not suppress its + // terminal notification or resource cleanup; the original rejection propagates. + if (killed) { + void evictTaskOutput(taskId) + // notified:true was pre-set so no XML notification fires; close the SDK + // task_started bookend directly. The in-process runner's own + // completion/failure emit guards on status==='running' so it won't + // double-emit after seeing status:killed. + emitTaskTerminatedSdk(taskId, 'stopped', { + toolUseId, + summary: description, + ownerAgentId: agentId ?? undefined, + }) + setTimeout( + evictTerminalTask.bind(null, taskId, setAppState), + STOPPED_DISPLAY_MS, + ) + } - if (killed) { - void evictTaskOutput(taskId) - // notified:true was pre-set so no XML notification fires; close the SDK - // task_started bookend directly. The in-process runner's own - // completion/failure emit guards on status==='running' so it won't - // double-emit after seeing status:killed. - emitTaskTerminatedSdk(taskId, 'stopped', { - toolUseId, - summary: description, - ownerAgentId: agentId ?? undefined, - }) - setTimeout( - evictTerminalTask.bind(null, taskId, setAppState), - STOPPED_DISPLAY_MS, - ) - } - - // Release perfetto agent registry entry - if (agentId) { - unregisterPerfettoAgent(agentId) + // Release perfetto agent registry entry + if (agentId) { + unregisterPerfettoAgent(agentId) + } } return killed