diff --git a/src/features/team-mode/team-state-store/store.test.ts b/src/features/team-mode/team-state-store/store.test.ts new file mode 100644 index 000000000..efc790447 --- /dev/null +++ b/src/features/team-mode/team-state-store/store.test.ts @@ -0,0 +1,269 @@ +/// + +import { afterEach, describe, expect, test } from "bun:test" +import { randomUUID } from "node:crypto" +import { mkdtemp, mkdir, readFile, rm, stat, utimes, writeFile } from "node:fs/promises" +import { tmpdir } from "node:os" +import path from "node:path" + +import { TeamModeConfigSchema } from "../../../config/schema/team-mode" +import type { TeamModeConfig } from "../../../config/schema/team-mode" +import type { RuntimeState, TeamSpec } from "../types" +import { + InvalidTransitionError, + RuntimeStateError, + STALE_DELETING_TTL_MS, + createRuntimeState, + listActiveTeams, + loadRuntimeState, + saveRuntimeState, + transitionRuntimeState, +} from "./store" + +async function createTemporaryBaseDir(): Promise { + return await mkdtemp(path.join(tmpdir(), "team-mode-store-")) +} + +function createConfig(baseDir: string): TeamModeConfig { + return TeamModeConfigSchema.parse({ + base_dir: baseDir, + max_members: 6, + max_parallel_members: 3, + max_messages_per_run: 200, + max_wall_clock_minutes: 45, + max_member_turns: 50, + }) +} + +function createSpec(name = `team-${randomUUID().slice(0, 8)}`): TeamSpec { + return { + version: 1, + name, + createdAt: Date.now(), + leadAgentId: "lead", + members: [ + { + kind: "subagent_type", + name: "lead", + subagent_type: "sisyphus", + backendType: "in-process", + isActive: true, + color: "red", + }, + { + kind: "category", + name: "worker", + category: "deep", + prompt: "implement task", + backendType: "in-process", + isActive: true, + color: "blue", + }, + ], + } +} + +async function seedRuntimeState( + runtimeState: RuntimeState, + config: TeamModeConfig, + saveRuntimeState: (runtimeState: RuntimeState, config: TeamModeConfig) => Promise, +): Promise { + await mkdir(path.join(config.base_dir ?? "", "runtime", runtimeState.teamRunId), { recursive: true }) + await saveRuntimeState(runtimeState, config) +} + +async function runtimeDirectoryExists(baseDir: string, teamRunId: string): Promise { + try { + await stat(path.join(baseDir, "runtime", teamRunId)) + return true + } catch (error) { + const nodeError = error as NodeJS.ErrnoException + if (nodeError.code === "ENOENT") return false + throw error + } +} + +describe("runtime state store", () => { + const temporaryDirectories: string[] = [] + + afterEach(async () => { + await Promise.all(temporaryDirectories.splice(0).map(async (directoryPath) => { + await rm(directoryPath, { recursive: true, force: true }) + })) + }) + + test("createRuntimeState persists creating state with computed bounds", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + + // when + const runtimeState = await createRuntimeState(createSpec(), undefined, "user", config) + const persistedState = JSON.parse(await readFile(path.join(baseDir, "runtime", runtimeState.teamRunId, "state.json"), "utf8")) + + // then + expect(runtimeState.teamRunId).toMatch(/^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i) + expect(runtimeState.status).toBe("creating") + expect(runtimeState.leadSessionId).toBeUndefined() + expect(runtimeState.bounds).toEqual({ + maxMembers: 6, + maxParallelMembers: 3, + maxMessagesPerRun: 200, + maxWallClockMinutes: 45, + maxMemberTurns: 50, + }) + expect(runtimeState.members).toEqual([ + expect.objectContaining({ name: "lead", agentType: "leader", status: "pending", pendingInjectedMessageIds: [] }), + expect.objectContaining({ name: "worker", agentType: "general-purpose", status: "pending", pendingInjectedMessageIds: [] }), + ]) + expect(persistedState.status).toBe("creating") + }) + + test("loadRuntimeState throws RuntimeStateError for malformed state", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const teamRunId = randomUUID() + await mkdir(path.join(baseDir, "runtime", teamRunId), { recursive: true }) + await writeFile(path.join(baseDir, "runtime", teamRunId, "state.json"), "{not-json") + + // when + const result = loadRuntimeState(teamRunId, config) + + // then + expect(result).rejects.toBeInstanceOf(RuntimeStateError) + }) + + test("transitionRuntimeState allows active to shutdown_requested", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const createdState = await createRuntimeState(createSpec(), "lead-session", "project", config) + + // when + await transitionRuntimeState(createdState.teamRunId, (runtimeState) => ({ ...runtimeState, status: "active" }), config) + const runtimeState = await transitionRuntimeState( + createdState.teamRunId, + (currentRuntimeState) => ({ ...currentRuntimeState, status: "shutdown_requested" }), + config, + ) + + // then + expect(runtimeState.status).toBe("shutdown_requested") + expect((await loadRuntimeState(createdState.teamRunId, config)).status).toBe("shutdown_requested") + }) + + test("transitionRuntimeState rejects reverse transition", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const createdState = await createRuntimeState(createSpec(), undefined, "user", config) + await seedRuntimeState({ ...createdState, status: "deleted" }, config, saveRuntimeState) + + // when + const result = transitionRuntimeState( + createdState.teamRunId, + (runtimeState) => ({ ...runtimeState, status: "active" }), + config, + ) + + // then + expect(result).rejects.toBeInstanceOf(InvalidTransitionError) + expect((await loadRuntimeState(createdState.teamRunId, config)).status).toBe("deleted") + }) + + test("loadRuntimeState ignores crash-left tmp files and keeps valid persisted state", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const runtimeState = await createRuntimeState(createSpec(), undefined, "user", config) + const statePath = path.join(baseDir, "runtime", runtimeState.teamRunId, "state.json") + await writeFile(`${statePath}.tmp.mock-crash`, JSON.stringify({ ...runtimeState, status: "active" })) + + // when + const persistedState = await loadRuntimeState(runtimeState.teamRunId, config) + + // then + expect(persistedState.status).toBe("creating") + }) + + test("loadRuntimeState accepts legacy member delegate counters without preserving them", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const runtimeState = await createRuntimeState(createSpec(), undefined, "user", config) + const statePath = path.join(baseDir, "runtime", runtimeState.teamRunId, "state.json") + await writeFile(statePath, JSON.stringify({ + ...runtimeState, + members: runtimeState.members.map((member) => ({ ...member, delegateTaskCallsUsed: 3 })), + })) + + // when + const persistedState = await loadRuntimeState(runtimeState.teamRunId, config) + + // then + expect(persistedState.members).toHaveLength(2) + expect(Object.keys(persistedState.members[0] ?? {})).not.toContain("delegateTaskCallsUsed") + }) + + test("listActiveTeams skips malformed runtime states and logs them", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const firstState = await createRuntimeState(createSpec("alpha-team"), undefined, "user", config) + const secondState = await createRuntimeState(createSpec("beta-team"), undefined, "project", config) + const malformedTeamRunId = randomUUID() + await mkdir(path.join(baseDir, "runtime", malformedTeamRunId), { recursive: true }) + await writeFile(path.join(baseDir, "runtime", malformedTeamRunId, "state.json"), "{oops") + + // when + const activeTeams = await listActiveTeams(config) + + // then + expect(activeTeams).toEqual([ + { teamRunId: firstState.teamRunId, teamName: "alpha-team", status: "creating", memberCount: 2, scope: "user" }, + { teamRunId: secondState.teamRunId, teamName: "beta-team", status: "creating", memberCount: 2, scope: "project" }, + ]) + }) + + test("listActiveTeams removes deleted runtime directories left by interrupted cleanup", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const runtimeState = await createRuntimeState(createSpec("deleted-team"), undefined, "user", config) + await saveRuntimeState({ ...runtimeState, status: "deleted" }, config) + + // when + const activeTeams = await listActiveTeams(config) + + // then + expect(activeTeams).toEqual([]) + expect(await runtimeDirectoryExists(baseDir, runtimeState.teamRunId)).toBe(false) + }) + + test("listActiveTeams removes deleting runtimes that have been stuck past the stale timeout", async () => { + // given + const baseDir = await createTemporaryBaseDir() + temporaryDirectories.push(baseDir) + const config = createConfig(baseDir) + const runtimeState = await createRuntimeState(createSpec("stuck-delete-team"), undefined, "user", config) + await saveRuntimeState({ ...runtimeState, status: "deleting" }, config) + const staleTimestamp = new Date(Date.now() - STALE_DELETING_TTL_MS - 1_000) + await utimes(path.join(baseDir, "runtime", runtimeState.teamRunId, "state.json"), staleTimestamp, staleTimestamp) + + // when + const activeTeams = await listActiveTeams(config) + + // then + expect(activeTeams).toEqual([]) + expect(await runtimeDirectoryExists(baseDir, runtimeState.teamRunId)).toBe(false) + }) +}) diff --git a/src/features/team-mode/team-state-store/store.ts b/src/features/team-mode/team-state-store/store.ts new file mode 100644 index 000000000..31999c3f1 --- /dev/null +++ b/src/features/team-mode/team-state-store/store.ts @@ -0,0 +1,253 @@ +import { randomUUID } from "node:crypto" +import { mkdir, readFile, readdir, rm, stat } from "node:fs/promises" +import path from "node:path" + +import type { TeamModeConfig } from "../../../config/schema/team-mode" +import { log } from "../../../shared/logger" +import { type RuntimeState, RuntimeStateSchema, type TeamSpec } from "../types" +import { getRuntimeStateDir, resolveBaseDir } from "../team-registry/paths" +import { atomicWrite, withLock } from "./locks" + +const STATE_FILE_NAME = "state.json" +export const STALE_DELETING_TTL_MS = 60_000 + +const ALLOWED_RUNTIME_TRANSITIONS: Readonly>> = { + creating: new Set(["active", "failed"]), + active: new Set(["shutdown_requested", "deleting"]), + shutdown_requested: new Set(["deleting"]), + deleting: new Set(["deleted"]), + deleted: new Set(), + failed: new Set(), + orphaned: new Set(), +} + +export class RuntimeStateError extends Error { + constructor(message: string, public readonly code: string) { + super(message) + this.name = "RuntimeStateError" + } +} + +export class InvalidTransitionError extends Error { + constructor(from: string, to: string) { + super(`invalid transition ${from} -> ${to}`) + this.name = "InvalidTransitionError" + } +} + +function getStatePath(baseDir: string, teamRunId: string): string { + return path.join(getRuntimeStateDir(baseDir, teamRunId), STATE_FILE_NAME) +} + +async function removeRuntimeDirectoryBestEffort( + baseDir: string, + teamRunId: string, + reason: "deleted" | "failed" | "stale_deleting", +): Promise { + try { + await rm(getRuntimeStateDir(baseDir, teamRunId), { recursive: true, force: true }) + } catch (error) { + log("team runtime cleanup failed", { + event: "team-runtime-cleanup-failed", + teamRunId, + reason, + error: error instanceof Error ? error.message : String(error), + }) + } +} + +async function isDeletingRuntimeStale(baseDir: string, teamRunId: string, now: number): Promise { + try { + const runtimeStateStat = await stat(getStatePath(baseDir, teamRunId)) + return now - runtimeStateStat.mtimeMs > STALE_DELETING_TTL_MS + } catch (error) { + const nodeError = error as NodeJS.ErrnoException + if (nodeError.code === "ENOENT") return true + throw error + } +} + +function serializeRuntimeState(runtimeState: RuntimeState): string { + const parsedRuntimeState = RuntimeStateSchema.parse(runtimeState) + return `${JSON.stringify(parsedRuntimeState, null, 2)}\n` +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function stripLegacyRuntimeStateMemberFields(member: unknown): unknown { + if (!isRecord(member)) { + return member + } + + const { delegateTaskCallsUsed: _delegateTaskCallsUsed, ...memberWithoutLegacyFields } = member + return memberWithoutLegacyFields +} + +function stripLegacyRuntimeStateFields(rawState: unknown): unknown { + if (!isRecord(rawState)) { + return rawState + } + + const members = rawState["members"] + if (!Array.isArray(members)) { + return rawState + } + + return { + ...rawState, + members: members.map(stripLegacyRuntimeStateMemberFields), + } +} + +function validateRuntimeState(rawState: unknown, teamRunId: string): RuntimeState { + const parsedRuntimeState = RuntimeStateSchema.safeParse(stripLegacyRuntimeStateFields(rawState)) + if (!parsedRuntimeState.success) { + throw new RuntimeStateError( + `runtime state invalid for ${teamRunId}: ${parsedRuntimeState.error.message}`, + "invalid_runtime_state", + ) + } + + return parsedRuntimeState.data +} + +function isValidTransition(fromStatus: RuntimeState["status"], toStatus: RuntimeState["status"]): boolean { + if (fromStatus === toStatus) return true + if (toStatus === "orphaned") return true + return ALLOWED_RUNTIME_TRANSITIONS[fromStatus].has(toStatus) +} + +export async function createRuntimeState( + spec: TeamSpec, + leadSessionId: string | undefined, + specSource: "project" | "user", + config: TeamModeConfig, +): Promise { + const baseDir = resolveBaseDir(config) + const teamRunId = randomUUID() + const runtimeDirectoryPath = getRuntimeStateDir(baseDir, teamRunId) + const runtimeState = validateRuntimeState({ + version: 1, + teamRunId, + teamName: spec.name, + specSource, + createdAt: Date.now(), + status: "creating", + leadSessionId, + members: spec.members.map((member) => ({ + name: member.name, + agentType: spec.leadAgentId === member.name ? "leader" : "general-purpose", + status: "pending", + color: member.color, + worktreePath: member.worktreePath, + lastInjectedTurnMarker: undefined, + pendingInjectedMessageIds: [], + })), + shutdownRequests: [], + bounds: { + maxMembers: config.max_members, + maxParallelMembers: config.max_parallel_members, + maxMessagesPerRun: config.max_messages_per_run, + maxWallClockMinutes: config.max_wall_clock_minutes, + maxMemberTurns: config.max_member_turns, + }, + }, teamRunId) + + await mkdir(runtimeDirectoryPath, { recursive: true }) + await atomicWrite(getStatePath(baseDir, teamRunId), serializeRuntimeState(runtimeState)) + return runtimeState +} + +export async function loadRuntimeState(teamRunId: string, config: TeamModeConfig): Promise { + const baseDir = resolveBaseDir(config) + const stateContent = await readFile(getStatePath(baseDir, teamRunId), "utf8") + + try { + return validateRuntimeState(JSON.parse(stateContent), teamRunId) + } catch (error) { + if (error instanceof RuntimeStateError) throw error + throw new RuntimeStateError( + `runtime state invalid for ${teamRunId}: ${(error as Error).message}`, + "invalid_runtime_state", + ) + } +} + +export async function saveRuntimeState(runtimeState: RuntimeState, config: TeamModeConfig): Promise { + const baseDir = resolveBaseDir(config) + await atomicWrite(getStatePath(baseDir, runtimeState.teamRunId), serializeRuntimeState(runtimeState)) +} + +export async function transitionRuntimeState( + teamRunId: string, + transition: (runtimeState: RuntimeState) => RuntimeState, + config: TeamModeConfig, +): Promise { + const baseDir = resolveBaseDir(config) + const runtimeDirectoryPath = getRuntimeStateDir(baseDir, teamRunId) + + return withLock(path.join(runtimeDirectoryPath, "state.lock"), async () => { + const currentRuntimeState = await loadRuntimeState(teamRunId, config) + const nextRuntimeState = validateRuntimeState(transition(currentRuntimeState), teamRunId) + + if (!isValidTransition(currentRuntimeState.status, nextRuntimeState.status)) { + throw new InvalidTransitionError(currentRuntimeState.status, nextRuntimeState.status) + } + + await saveRuntimeState(nextRuntimeState, config) + return nextRuntimeState + }, { ownerTag: "team-state-store" }) +} + +export async function listActiveTeams( + config: TeamModeConfig, +): Promise> { + const baseDir = resolveBaseDir(config) + const now = Date.now() + + try { + const runtimeEntries = await readdir(path.join(baseDir, "runtime"), { withFileTypes: true }) + const activeTeams: Array<{ teamRunId: string; teamName: string; status: string; memberCount: number; scope: "project" | "user" }> = [] + + for (const runtimeEntry of runtimeEntries) { + if (!runtimeEntry.isDirectory()) continue + + try { + const runtimeState = await loadRuntimeState(runtimeEntry.name, config) + + if (runtimeState.status === "deleted" || runtimeState.status === "failed") { + await removeRuntimeDirectoryBestEffort(baseDir, runtimeEntry.name, runtimeState.status) + continue + } + + if (runtimeState.status === "deleting" && await isDeletingRuntimeStale(baseDir, runtimeEntry.name, now)) { + await removeRuntimeDirectoryBestEffort(baseDir, runtimeEntry.name, "stale_deleting") + continue + } + + activeTeams.push({ + teamRunId: runtimeState.teamRunId, + teamName: runtimeState.teamName, + status: runtimeState.status, + memberCount: runtimeState.members.length, + scope: runtimeState.specSource, + }) + } catch (error) { + log("team runtime state skipped", { + event: "team-runtime-state-skipped", + teamRunId: runtimeEntry.name, + error: error instanceof Error ? error.message : String(error), + }) + } + } + + activeTeams.sort((leftTeam, rightTeam) => leftTeam.teamName.localeCompare(rightTeam.teamName) || leftTeam.teamRunId.localeCompare(rightTeam.teamRunId)) + return activeTeams + } catch (error) { + const nodeError = error as NodeJS.ErrnoException + if (nodeError.code === "ENOENT") return [] + throw error + } +}