From 8af33fc9e0fe42cc125d4cf40b30d18194913618 Mon Sep 17 00:00:00 2001 From: Gabriel Brown Date: Fri, 10 Jul 2026 17:33:38 -0400 Subject: [PATCH] fix(worker): always release stopped workspaces --- apps/agent-worker/src/worker.ts | 43 +++- .../maintenance-workspace-cleanup.test.ts | 187 ++++++++++++++++++ 2 files changed, 222 insertions(+), 8 deletions(-) create mode 100644 apps/agent-worker/tests/unit/maintenance-workspace-cleanup.test.ts diff --git a/apps/agent-worker/src/worker.ts b/apps/agent-worker/src/worker.ts index 5426d6f..3c6bc3b 100644 --- a/apps/agent-worker/src/worker.ts +++ b/apps/agent-worker/src/worker.ts @@ -139,6 +139,13 @@ const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); const client = new ConvexHttpClient(env.convexUrl); const activeWorkspaces = new Map(); +export const _setActiveWorkspaceForTests = ( + jobId: string, + workspace: ActiveWorkspace, +) => activeWorkspaces.set(jobId, workspace); + +export const _resetActiveWorkspacesForTests = () => activeWorkspaces.clear(); + const appendEvent = async ( jobId: Id<'agentJobs'>, level: 'debug' | 'info' | 'warn' | 'error', @@ -1894,15 +1901,35 @@ export const openWorkspacePullRequest = async (jobId: string) => { export const stopWorkspace = async (jobId: string) => { const workspace = resolveWorkspace(jobId); - await markWorkspaceStopped(workspace.claim.job._id); - workspace.opencodeSession?.close(); - if (workspace.containerName) { - await stopWorkspaceContainer(workspace.containerName); + const errors: unknown[] = []; + try { + await markWorkspaceStopped(workspace.claim.job._id); + } catch (error) { + errors.push(error); + } + try { + workspace.opencodeSession?.close(); + } catch (error) { + errors.push(error); + } + try { + if (workspace.containerName) { + await stopWorkspaceContainer(workspace.containerName); + } + } catch (error) { + errors.push(error); + } finally { + activeWorkspaces.delete(jobId); + try { + workspace.boxHandle.release(); + } catch (error) { + errors.push(error); + } + } + if (errors.length === 1) throw errors[0]; + if (errors.length > 1) { + throw new AggregateError(errors, 'Workspace teardown failed.'); } - activeWorkspaces.delete(jobId); - // The persistent per-user home + ~/Code checkouts survive across sessions; - // release the box (reaped once no other thread/terminal holds it). - workspace.boxHandle.release(); return { success: true }; }; diff --git a/apps/agent-worker/tests/unit/maintenance-workspace-cleanup.test.ts b/apps/agent-worker/tests/unit/maintenance-workspace-cleanup.test.ts new file mode 100644 index 0000000..5e7c6b4 --- /dev/null +++ b/apps/agent-worker/tests/unit/maintenance-workspace-cleanup.test.ts @@ -0,0 +1,187 @@ +import { getFunctionName } from 'convex/server'; +import { afterEach, describe, expect, test, vi } from 'vitest'; + +import type { Id } from '@spoon/backend/convex/_generated/dataModel.js'; + +const mocks = vi.hoisted(() => ({ + getWorktreeDiff: vi.fn(() => + Promise.resolve({ exitCode: 0, output: 'unexpected diff' }), + ), + mutation: vi.fn(), + stopWorkspaceContainer: vi.fn(() => Promise.resolve()), + streamExecInContainer: vi.fn(), +})); + +vi.mock('convex/browser', () => ({ + ConvexHttpClient: vi.fn( + class ConvexHttpClient { + mutation = mocks.mutation; + }, + ), +})); +vi.mock('../../src/env', () => ({ + env: { + buildCreatedAt: 'test', + buildSha: 'test', + containerAccess: 'network', + convexUrl: 'http://convex.test', + jobTimeoutMs: 1000, + runtime: 'docker', + workerId: 'worker-1', + workerToken: 'token', + }, +})); +vi.mock('../../src/git', () => ({ + cloneRepository: vi.fn(), + commitAndPush: vi.fn(), + getStatus: vi.fn(), + getWorktreeDiff: mocks.getWorktreeDiff, + run: vi.fn(), +})); +vi.mock('../../src/github', () => ({ + getInstallationToken: vi.fn(), + openDraftPullRequest: vi.fn(), +})); +vi.mock('../../src/opencode-session', () => ({ + abortOpenCodeSession: vi.fn(), + createOpenCodeSession: vi.fn(), + promptOpenCodeSession: vi.fn(), + replyOpenCodePermission: vi.fn(), +})); +vi.mock('../../src/runtime/docker', () => ({ + listWorkspaceContainerNames: vi.fn(() => Promise.resolve([])), + runExecInContainer: vi.fn(), + startWorkspaceContainer: vi.fn(), + stopWorkspaceContainer: mocks.stopWorkspaceContainer, + streamExecInContainer: mocks.streamExecInContainer, +})); +vi.mock('../../src/user-container', () => ({ + acquireUserBox: vi.fn(), + runningBoxUsernames: vi.fn(() => new Set()), +})); +vi.mock('../../src/user-environment', () => ({ + fetchUserEnvironment: vi.fn(), + materializeUserHome: vi.fn(), +})); + +const decision = { + decision: 'sync', + risk: 'low', + summary: 'Safe to sync.', + ignoredCommitShas: [], + ignoredReason: '', + recommendedAction: 'Sync upstream.', + requiresUserApproval: false, +}; + +const load = async () => await import('../../src/worker'); + +afterEach(async () => { + const worker = await load(); + worker._resetActiveWorkspacesForTests(); + vi.restoreAllMocks(); + vi.clearAllMocks(); +}); + +describe('maintenance workspace cleanup', () => { + test.each(['mark', 'container'] as const)( + 'releases and removes the workspace when %s teardown fails', + async (failingStep) => { + const worker = await load(); + const release = vi.fn(); + const close = vi.fn(); + const jobId = `maintenance-${failingStep}`; + mocks.mutation.mockImplementation((reference) => { + if ( + failingStep === 'mark' && + getFunctionName(reference) === 'agentJobs:markWorkspaceStopped' + ) { + return Promise.reject(new Error('status teardown failed')); + } + return Promise.resolve('message-1'); + }); + mocks.stopWorkspaceContainer.mockImplementation(() => + failingStep === 'container' + ? Promise.reject(new Error('container teardown failed')) + : Promise.resolve(), + ); + mocks.streamExecInContainer.mockImplementation(async (args) => { + await args.onStdoutLine( + JSON.stringify({ + type: 'item.completed', + item: { + type: 'agent_message', + text: JSON.stringify(decision), + }, + }), + ); + return { exitCode: 0, output: '' }; + }); + worker._setActiveWorkspaceForTests(jobId, { + claim: { + job: { + _id: jobId as Id<'agentJobs'>, + prompt: 'Review maintenance changes', + jobType: 'maintenance_review', + baseBranch: 'main', + workBranch: 'spoon/maintenance', + forkOwner: 'team', + forkRepo: 'spoon', + upstreamOwner: 'upstream', + upstreamRepo: 'spoon', + }, + spoon: { name: 'Spoon' }, + openai: { + model: 'gpt-5', + reasoningEffort: 'medium', + }, + aiProviderProfile: { + id: 'profile-1', + name: 'Codex', + provider: 'opencode_openai_login', + authType: 'opencode_auth_json', + model: 'gpt-5', + reasoningEffort: 'medium', + }, + github: {}, + secrets: [], + }, + workdir: '/tmp/workspace', + homeDir: '/tmp/workspace', + username: 'tester', + containerHome: '/home/tester', + containerRepo: '/home/tester/Code/spoon', + repoDir: '/tmp/workspace/Code/spoon', + boxName: 'spoon-box-tester', + boxHandle: { boxName: 'spoon-box-tester', release }, + githubToken: 'github-token', + redact: (value: string) => value, + containerName: 'spoon-agent-job-maintenance', + opencodeSession: { + client: {} as never, + sessionId: 'session-1', + close, + }, + }); + const consoleError = vi + .spyOn(console, 'error') + .mockImplementation(() => undefined); + + await expect( + worker.sendWorkspaceMessage(jobId, 'Review maintenance changes'), + ).resolves.toBeUndefined(); + + expect(release).toHaveBeenCalledTimes(1); + expect(close).toHaveBeenCalledTimes(1); + expect(mocks.stopWorkspaceContainer).toHaveBeenCalledTimes(1); + expect(mocks.getWorktreeDiff).not.toHaveBeenCalled(); + expect( + mocks.mutation.mock.calls.some(([, args]) => args.kind === 'diff'), + ).toBe(false); + await expect(worker.stopWorkspace(jobId)).rejects.toThrow( + 'Agent workspace is not active on this worker.', + ); + expect(consoleError).toHaveBeenCalledTimes(1); + }, + ); +});