From daec4d6dfa45e33f716edc8ca3d4111cee16ed3c Mon Sep 17 00:00:00 2001 From: Gabriel Brown Date: Fri, 10 Jul 2026 15:55:11 -0400 Subject: [PATCH] fix(worker): make workspace cleanup layout-aware --- apps/agent-worker/src/user-container.ts | 2 + apps/agent-worker/src/worker.ts | 95 ++++++++++++++++--- .../tests/unit/cleanup-layout.test.ts | 73 ++++++++++++++ 3 files changed, 156 insertions(+), 14 deletions(-) create mode 100644 apps/agent-worker/tests/unit/cleanup-layout.test.ts diff --git a/apps/agent-worker/src/user-container.ts b/apps/agent-worker/src/user-container.ts index c4bb374..5f52764 100644 --- a/apps/agent-worker/src/user-container.ts +++ b/apps/agent-worker/src/user-container.ts @@ -38,3 +38,5 @@ export const releaseUserBox = (username: string) => { boxes.delete(username); }, env.boxIdleMs); }; + +export const runningBoxUsernames = (): Set => new Set(boxes.keys()); diff --git a/apps/agent-worker/src/worker.ts b/apps/agent-worker/src/worker.ts index d5025ad..831cf2a 100644 --- a/apps/agent-worker/src/worker.ts +++ b/apps/agent-worker/src/worker.ts @@ -41,7 +41,11 @@ import { stopWorkspaceContainer, streamExecInContainer, } from './runtime/docker'; -import { acquireUserBox, releaseUserBox } from './user-container'; +import { + acquireUserBox, + releaseUserBox, + runningBoxUsernames, +} from './user-container'; import { fetchUserEnvironment, materializeUserHome } from './user-environment'; type Claim = { @@ -1287,6 +1291,29 @@ const slugify = (value: string) => .replace(/^-+|-+$/g, '') .replace(/-{2,}/g, '-') || 'x'; +export const planWorkdirCleanup = (args: { + root: string; + rootEntries: { name: string; isDirectory: boolean }[]; + homesEntries: Record; + codeLeaves: Record; + activeWorkdirs: Set; + runningBoxUsernames: Set; +}): { removeDirs: string[] } => { + const removeDirs: string[] = []; + for (const entry of args.rootEntries) { + if (!entry.isDirectory || entry.name.startsWith('.')) continue; + if (entry.name === 'homes') continue; + const abs = path.resolve(args.root, entry.name); + if (args.activeWorkdirs.has(abs)) continue; + removeDirs.push(abs); + } + for (const [username, leaves] of Object.entries(args.codeLeaves)) { + if (args.runningBoxUsernames.has(username)) continue; + for (const leaf of leaves) removeDirs.push(leaf); + } + return { removeDirs }; +}; + const runClaim = async (claim: Claim) => { const jobId = claim.job._id; const secretValues = [ @@ -1885,7 +1912,8 @@ export const getWorkerHealth = async () => { workdir: workspace.workdir, agentTurnActive: Boolean(workspace.agentTurnActive), })); - const containerNames = await listWorkspaceContainerNames('spoon-agent-job-'); + const jobContainers = await listWorkspaceContainerNames('spoon-agent-job-'); + const boxContainers = await listWorkspaceContainerNames('spoon-box-'); return { ok: true, buildSha: env.buildSha, @@ -1904,7 +1932,8 @@ export const getWorkerHealth = async () => { jobTimeoutMs: env.jobTimeoutMs, activeWorkspaceCount: active.length, activeWorkspaces: active, - workspaceContainers: containerNames, + workspaceContainers: jobContainers, + boxContainers, }; }; @@ -1928,14 +1957,55 @@ export const cleanupOrphanedWorkspaces = async () => { removedContainers.push(containerName); } - const removedWorkdirs: string[] = []; const root = path.resolve(env.workdir); + const removedWorkdirs: string[] = []; try { - const entries = await readdir(root, { withFileTypes: true }); - for (const entry of entries) { - if (!entry.isDirectory() || entry.name.startsWith('.')) continue; - const target = path.resolve(root, entry.name); - if (activeWorkdirs.has(target)) continue; + const rootDirents = await readdir(root, { withFileTypes: true }); + const rootEntries = rootDirents.map((d) => ({ + name: d.name, + isDirectory: d.isDirectory(), + })); + const homesRoot = path.join(root, 'homes'); + const homesEntries: Record< + string, + { name: string; isDirectory: boolean }[] + > = {}; + const codeLeaves: Record = {}; + const homesDirents = await readdir(homesRoot, { + withFileTypes: true, + }).catch(() => []); + for (const userDir of homesDirents) { + if (!userDir.isDirectory()) continue; + const username = userDir.name; + const codeRoot = path.join(homesRoot, username, 'Code'); + homesEntries[username] = []; + const leaves: string[] = []; + const spoons = await readdir(codeRoot, { withFileTypes: true }).catch( + () => [], + ); + for (const spoonDir of spoons) { + if (!spoonDir.isDirectory()) continue; + const spoonRoot = path.join(codeRoot, spoonDir.name); + const branches = await readdir(spoonRoot, { + withFileTypes: true, + }).catch(() => []); + for (const branchDir of branches) { + if (branchDir.isDirectory()) { + leaves.push(path.join(spoonRoot, branchDir.name)); + } + } + } + codeLeaves[username] = leaves; + } + const { removeDirs } = planWorkdirCleanup({ + root, + rootEntries, + homesEntries, + codeLeaves, + activeWorkdirs, + runningBoxUsernames: runningBoxUsernames(), + }); + for (const target of removeDirs) { await rm(target, { recursive: true, force: true }); removedWorkdirs.push(target); } @@ -1946,11 +2016,8 @@ export const cleanupOrphanedWorkspaces = async () => { } } - return { - success: true, - removedContainers, - removedWorkdirs, - }; + const boxContainers = await listWorkspaceContainerNames('spoon-box-'); + return { success: true, removedContainers, removedWorkdirs, boxContainers }; }; export const startWorker = async () => { diff --git a/apps/agent-worker/tests/unit/cleanup-layout.test.ts b/apps/agent-worker/tests/unit/cleanup-layout.test.ts new file mode 100644 index 0000000..d30053c --- /dev/null +++ b/apps/agent-worker/tests/unit/cleanup-layout.test.ts @@ -0,0 +1,73 @@ +import { afterEach, describe, expect, test, vi } from 'vitest'; + +const load = async () => { + vi.resetModules(); + process.env.SPOON_WORKER_TOKEN = 'test-worker-token'; + process.env.GITHUB_APP_ID = '123'; + process.env.GITHUB_APP_PRIVATE_KEY = + '-----BEGIN PRIVATE KEY-----\ntest\n-----END PRIVATE KEY-----'; + return await import('../../src/worker'); +}; + +describe('planWorkdirCleanup', () => { + afterEach(() => vi.resetModules()); + + test('never removes the homes/ root or any home directory', async () => { + const { planWorkdirCleanup } = await load(); + const plan = planWorkdirCleanup({ + root: '/work', + rootEntries: [{ name: 'homes', isDirectory: true }], + homesEntries: { alice: [{ name: 'Code', isDirectory: true }] }, + codeLeaves: {}, + activeWorkdirs: new Set(), + runningBoxUsernames: new Set(), + }); + expect(plan.removeDirs).not.toContain('/work/homes'); + expect(plan.removeDirs).not.toContain('/work/homes/alice'); + }); + + test('removes legacy top-level job dirs not in the active set', async () => { + const { planWorkdirCleanup } = await load(); + const plan = planWorkdirCleanup({ + root: '/work', + rootEntries: [ + { name: 'homes', isDirectory: true }, + { name: 'legacy-job-123', isDirectory: true }, + { name: 'dev', isDirectory: true }, + ], + homesEntries: {}, + codeLeaves: {}, + activeWorkdirs: new Set(['/work/dev']), + runningBoxUsernames: new Set(), + }); + expect(plan.removeDirs).toContain('/work/legacy-job-123'); + expect(plan.removeDirs).not.toContain('/work/dev'); + expect(plan.removeDirs).not.toContain('/work/homes'); + }); + + test('removes per-thread checkouts only for users with no running box', async () => { + const { planWorkdirCleanup } = await load(); + const plan = planWorkdirCleanup({ + root: '/work', + rootEntries: [{ name: 'homes', isDirectory: true }], + homesEntries: { + alice: [{ name: 'Code', isDirectory: true }], + bob: [{ name: 'Code', isDirectory: true }], + }, + codeLeaves: { + alice: ['/work/homes/alice/Code/spoon-a/branch-x'], + bob: ['/work/homes/bob/Code/spoon-b/branch-y'], + }, + activeWorkdirs: new Set(), + runningBoxUsernames: new Set(['bob']), + }); + expect(plan.removeDirs).toContain( + '/work/homes/alice/Code/spoon-a/branch-x', + ); + expect(plan.removeDirs).not.toContain( + '/work/homes/bob/Code/spoon-b/branch-y', + ); + expect(plan.removeDirs).not.toContain('/work/homes/alice'); + expect(plan.removeDirs).not.toContain('/work/homes/alice/Code'); + }); +});