From cb1c7508e7819cd8e0812838cd58fa28cc7f63fb Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Tue, 28 Apr 2026 10:46:54 +0900 Subject: [PATCH] feat(team-mode): add team runtime resource cleanup with tests --- .../cleanup-team-run-resources.test.ts | 80 +++++++++++++++++++ .../cleanup-team-run-resources.ts | 77 ++++++++++++++++++ 2 files changed, 157 insertions(+) create mode 100644 src/features/team-mode/team-runtime/cleanup-team-run-resources.test.ts create mode 100644 src/features/team-mode/team-runtime/cleanup-team-run-resources.ts diff --git a/src/features/team-mode/team-runtime/cleanup-team-run-resources.test.ts b/src/features/team-mode/team-runtime/cleanup-team-run-resources.test.ts new file mode 100644 index 000000000..22ee7d53d --- /dev/null +++ b/src/features/team-mode/team-runtime/cleanup-team-run-resources.test.ts @@ -0,0 +1,80 @@ +/// + +import { afterEach, describe, expect, test } from "bun:test" +import { mkdir, mkdtemp, rm } from "node:fs/promises" +import { tmpdir } from "node:os" +import path from "node:path" + +import type { TeamModeConfig } from "../../../config/schema/team-mode" +import { TeamModeConfigSchema } from "../../../config/schema/team-mode" +import type { BackgroundManager } from "../../background-agent/manager" +import { + clearTeamSessionRegistry, + lookupTeamSession, + registerTeamSession, +} from "../team-session-registry" +import { saveRuntimeState } from "../team-state-store/store" +import type { RuntimeState } from "../types" +import { cleanupTeamRunResources } from "./cleanup-team-run-resources" + +const temporaryDirectories: string[] = [] + +function createConfig(baseDir: string): TeamModeConfig { + return TeamModeConfigSchema.parse({ base_dir: baseDir, enabled: true }) +} + +function createRuntimeState(teamRunId: string): RuntimeState { + return { + version: 1, + teamRunId, + teamName: "team-alpha", + specSource: "project", + createdAt: 1, + status: "creating", + leadSessionId: "lead-session", + members: [ + { name: "worker-1", agentType: "general-purpose", status: "pending", pendingInjectedMessageIds: [] }, + ], + shutdownRequests: [], + bounds: { maxMembers: 8, maxParallelMembers: 4, maxMessagesPerRun: 10_000, maxWallClockMinutes: 120, maxMemberTurns: 500 }, + } +} + +function createStubBgMgr(): BackgroundManager { + return { + cancelTask: async () => undefined, + } as unknown as BackgroundManager +} + +describe("cleanupTeamRunResources", () => { + afterEach(async () => { + clearTeamSessionRegistry() + await Promise.all(temporaryDirectories.splice(0).map(async (directoryPath) => rm(directoryPath, { recursive: true, force: true }))) + }) + + test("unregisters every team-session-registry entry for the failed team so the gating hook cannot authorize stale participants", async () => { + // given + const baseDir = await mkdtemp(path.join(tmpdir(), "cleanup-team-run-registry-")) + temporaryDirectories.push(baseDir) + const teamRunId = "33333333-3333-4333-8333-333333333333" + await mkdir(path.join(baseDir, "runtime", teamRunId), { recursive: true }) + await saveRuntimeState(createRuntimeState(teamRunId), createConfig(baseDir)) + registerTeamSession("lead-session", { teamRunId, memberName: "lead", role: "lead" }) + registerTeamSession("worker-session", { teamRunId, memberName: "worker-1", role: "member" }) + registerTeamSession("other-team-session", { teamRunId: "other-team", memberName: "solo", role: "member" }) + + // when + await cleanupTeamRunResources({ + teamRunId, + config: createConfig(baseDir), + resources: [{}], + bgMgr: createStubBgMgr(), + createdLayout: false, + }) + + // then + expect(lookupTeamSession("lead-session")).toBeUndefined() + expect(lookupTeamSession("worker-session")).toBeUndefined() + expect(lookupTeamSession("other-team-session")).toEqual({ teamRunId: "other-team", memberName: "solo", role: "member" }) + }) +}) diff --git a/src/features/team-mode/team-runtime/cleanup-team-run-resources.ts b/src/features/team-mode/team-runtime/cleanup-team-run-resources.ts new file mode 100644 index 000000000..931339d4d --- /dev/null +++ b/src/features/team-mode/team-runtime/cleanup-team-run-resources.ts @@ -0,0 +1,77 @@ +import { rm } from "node:fs/promises" + +import type { TeamModeConfig } from "../../../config/schema/team-mode" +import type { BackgroundManager } from "../../background-agent/manager" +import type { TmuxSessionManager } from "../../tmux-subagent/manager" +import { removeTeamLayout } from "../team-layout-tmux/layout" +import { unregisterTeamSessionsByTeam } from "../team-session-registry" +import { loadRuntimeState, transitionRuntimeState } from "../team-state-store/store" +import type { TeamRunCreateError } from "./create" + +type SpawnedMemberResource = { + taskId?: string + worktreePath?: string +} + +function normalizeError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)) +} + +export async function cleanupTeamRunResources(args: { + teamRunId: string + config: TeamModeConfig + resources: SpawnedMemberResource[] + bgMgr: BackgroundManager + tmuxMgr?: TmuxSessionManager + createdLayout: boolean +}): Promise { + const cleanupReport: TeamRunCreateError["cleanupReport"] = { + cancelledTaskIds: [], + removedLayout: false, + removedWorktrees: [], + errors: [], + } + + for (const resource of [...args.resources].reverse()) { + if (resource.taskId) { + try { + await args.bgMgr.cancelTask(resource.taskId, { + source: "team-create-rollback", + reason: "creating_rollback", + skipNotification: true, + }) + cleanupReport.cancelledTaskIds.push(resource.taskId) + } catch (cancelError) { + cleanupReport.errors.push(`cancel ${resource.taskId}: ${normalizeError(cancelError).message}`) + } + } + + if (resource.worktreePath) { + try { + await rm(resource.worktreePath, { recursive: true, force: true }) + cleanupReport.removedWorktrees.push(resource.worktreePath) + } catch (cleanupError) { + cleanupReport.errors.push(`worktree ${resource.worktreePath}: ${normalizeError(cleanupError).message}`) + } + } + } + + if (args.createdLayout && args.tmuxMgr) { + try { + const runtimeState = await loadRuntimeState(args.teamRunId, args.config) + await removeTeamLayout(args.teamRunId, runtimeState.tmuxLayout, args.tmuxMgr) + cleanupReport.removedLayout = true + } catch (layoutError) { + cleanupReport.errors.push(`layout ${args.teamRunId}: ${normalizeError(layoutError).message}`) + } + } + + await transitionRuntimeState(args.teamRunId, (runtimeState) => ({ ...runtimeState, status: "failed" }), args.config).catch((transitionError) => { + cleanupReport.errors.push(`state ${args.teamRunId}: ${normalizeError(transitionError).message}`) + return undefined + }) + + unregisterTeamSessionsByTeam(args.teamRunId) + + return cleanupReport +}