feat(team-mode): add team runtime resource cleanup with tests

2026-04-28 10:46:54 +09:00
parent 62d409085e
commit cb1c7508e7
2 changed files with 157 additions and 0 deletions
@@ -0,0 +1,80 @@
+/// <reference types="bun-types" />
+
+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" })
+  })
+})
@@ -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<TeamRunCreateError["cleanupReport"]> {
+  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
+}