diff --git a/apps/agent-worker/src/index.ts b/apps/agent-worker/src/index.ts index 2237756..14d5841 100644 --- a/apps/agent-worker/src/index.ts +++ b/apps/agent-worker/src/index.ts @@ -1,5 +1,6 @@ import { env } from './env'; import { startWorkerServer } from './server'; +import { reconcileExistingBoxes } from './user-container'; import { startWorker } from './worker'; // Dev-only watchdog: the dev runner chain (turbo → with-env → dotenv → bash) @@ -26,4 +27,5 @@ if (env.devWatchdog) { } startWorkerServer(); +await reconcileExistingBoxes(); await startWorker(); diff --git a/apps/agent-worker/src/runtime/docker.ts b/apps/agent-worker/src/runtime/docker.ts index 9d79176..f169928 100644 --- a/apps/agent-worker/src/runtime/docker.ts +++ b/apps/agent-worker/src/runtime/docker.ts @@ -362,6 +362,7 @@ export const ensureUserContainer = async (args: { [ 'run', '-d', + '--init', '--name', name, '--memory', diff --git a/apps/agent-worker/src/terminal.ts b/apps/agent-worker/src/terminal.ts index 1a8d3be..c99a5a3 100644 --- a/apps/agent-worker/src/terminal.ts +++ b/apps/agent-worker/src/terminal.ts @@ -4,9 +4,10 @@ import type { Server } from 'node:http'; import type { WebSocket } from 'ws'; import { WebSocketServer } from 'ws'; +import type { BoxHandle } from './user-container'; import { env } from './env'; import { verifyTerminalToken } from './terminal-token'; -import { acquireUserBox, releaseUserBox } from './user-container'; +import { acquireUserBox } from './user-container'; import { getTerminalWorkspace } from './worker'; const clampDimension = (value: unknown) => { @@ -63,7 +64,7 @@ const bridge = async (ws: WebSocket, jobId: string) => { else pendingInput.push(data); }); - let acquired = false; + const handleHolder: { current?: BoxHandle } = {}; let released = false; // Read through a function so TS doesn't narrow `released` to a constant — the // cleanup handler flips it asynchronously when the socket closes. @@ -72,7 +73,7 @@ const bridge = async (ws: WebSocket, jobId: string) => { if (released) return; released = true; procHolder.current?.kill(); - if (acquired) releaseUserBox(workspace.username); + handleHolder.current?.release(); }; ws.on('close', cleanup); ws.on('error', cleanup); @@ -81,12 +82,13 @@ const bridge = async (ws: WebSocket, jobId: string) => { // the terminal share the exact same container (Phase 2). let boxName: string; try { - boxName = await acquireUserBox({ + const handle = await acquireUserBox({ username: workspace.username, workdir: workspace.workdir, containerHome: workspace.containerHome, }); - acquired = true; + handleHolder.current = handle; + boxName = handle.boxName; } catch (error) { ws.close( 1011, @@ -95,7 +97,11 @@ const bridge = async (ws: WebSocket, jobId: string) => { return; } - if (isReleased()) return; // client disconnected during startup; cleanup ran + // Cleanup may have run before the awaited handle existed. + if (isReleased()) { + handleHolder.current.release(); + return; + } // Reattach a persistent tmux session across reconnects when available, else a // plain login shell. `stty` sizes the PTY to the client's viewport up front. diff --git a/apps/agent-worker/src/user-container.ts b/apps/agent-worker/src/user-container.ts index 5f52764..4f119a2 100644 --- a/apps/agent-worker/src/user-container.ts +++ b/apps/agent-worker/src/user-container.ts @@ -1,42 +1,110 @@ import { env } from './env'; import { ensureUserContainer, + listWorkspaceContainerNames, stopWorkspaceContainer, userContainerName, } from './runtime/docker'; -// Phase 2: one persistent "box" container per user that all of their threads -// (agent turns + terminal + commands) exec into. Reference-counted so it stays -// up while any thread workspace is active or a terminal is connected, and is -// reaped after an idle period once nothing holds it. -type Box = { refs: number; idleTimer?: NodeJS.Timeout }; -const boxes = new Map(); +// Phase 2: one persistent "box" container per user. Reference-counted by opaque +// handles (not a shared integer), serialized per-username with an async mutex, +// and idle-reaped once no handle is held. +export type BoxHandle = { boxName: string; release: () => void }; -export const acquireUserBox = async (args: { - username: string; - workdir: string; - containerHome: string; -}): Promise => { - const name = await ensureUserContainer(args); - const box = boxes.get(args.username) ?? { refs: 0 }; - if (box.idleTimer) { - clearTimeout(box.idleTimer); - box.idleTimer = undefined; - } - box.refs += 1; - boxes.set(args.username, box); - return name; +type Box = { + name: string; + refs: Set; + idleTimer?: ReturnType; +}; +const boxes = new Map(); +const locks = new Map>(); + +// Per-username async mutex: chain each operation after the previous one. +const withLock = (username: string, fn: () => Promise): Promise => { + const previous = locks.get(username) ?? Promise.resolve(); + const next = previous.then(fn, fn); + locks.set( + username, + next.then( + () => undefined, + () => undefined, + ), + ); + return next; }; -export const releaseUserBox = (username: string) => { +const scheduleReapIfIdle = (username: string) => { const box = boxes.get(username); - if (!box) return; - box.refs = Math.max(0, box.refs - 1); - if (box.refs > 0) return; + if (!box || box.refs.size > 0) return; + if (box.idleTimer) clearTimeout(box.idleTimer); box.idleTimer = setTimeout(() => { void stopWorkspaceContainer(userContainerName(username)); boxes.delete(username); }, env.boxIdleMs); }; +const makeHandle = (username: string, box: Box): BoxHandle => { + const token = Symbol('box-ref'); + box.refs.add(token); + let released = false; + return { + boxName: box.name, + release: () => { + if (released) return; + released = true; + box.refs.delete(token); + scheduleReapIfIdle(username); + }, + }; +}; + +export const acquireUserBox = (args: { + username: string; + workdir: string; + containerHome: string; +}): Promise => + withLock(args.username, async () => { + let box = boxes.get(args.username); + let needsEnsure = false; + if (box?.idleTimer) { + clearTimeout(box.idleTimer); + box.idleTimer = undefined; + } + if (!box) { + box = { name: userContainerName(args.username), refs: new Set() }; + boxes.set(args.username, box); + needsEnsure = true; + } + // Register the ref before the slow Docker call so a failed acquire can + // release it without leaking registry state. + const handle = makeHandle(args.username, box); + try { + if (needsEnsure) box.name = await ensureUserContainer(args); + return handle; + } catch (error) { + handle.release(); + throw error; + } + }); + +// Adopt boxes left by a prior worker process so the idle reaper cleans them up. +export const reconcileExistingBoxes = async (): Promise => { + const names = await listWorkspaceContainerNames('spoon-box-'); + for (const name of names) { + const username = name.replace(/^spoon-box-/, ''); + if (boxes.has(username)) continue; + const box: Box = { name, refs: new Set() }; + boxes.set(username, box); + scheduleReapIfIdle(username); + } +}; + export const runningBoxUsernames = (): Set => new Set(boxes.keys()); + +export const _resetBoxRegistryForTests = () => { + for (const box of boxes.values()) { + if (box.idleTimer) clearTimeout(box.idleTimer); + } + boxes.clear(); + locks.clear(); +}; diff --git a/apps/agent-worker/src/worker.ts b/apps/agent-worker/src/worker.ts index 831cf2a..55b794a 100644 --- a/apps/agent-worker/src/worker.ts +++ b/apps/agent-worker/src/worker.ts @@ -16,6 +16,7 @@ import { api } from '@spoon/backend/convex/_generated/api.js'; import type { NormalizedAgentEvent } from './agent-events'; import type { OpenCodeSession } from './opencode-session'; +import type { BoxHandle } from './user-container'; import { normalizeCodexJsonLine } from './agent-events'; import { prepareCodexWorkspaceFiles } from './codex-runtime'; import { env } from './env'; @@ -41,11 +42,7 @@ import { stopWorkspaceContainer, streamExecInContainer, } from './runtime/docker'; -import { - acquireUserBox, - releaseUserBox, - runningBoxUsernames, -} from './user-container'; +import { acquireUserBox, runningBoxUsernames } from './user-container'; import { fetchUserEnvironment, materializeUserHome } from './user-environment'; type Claim = { @@ -111,6 +108,7 @@ type ActiveWorkspace = { repoDir: string; // Phase 2: the per-user box container this thread execs into. boxName: string; + boxHandle: BoxHandle; githubToken: string; redact: (value: string) => string; runtimeMode?: 'opencode_server' | 'codex_exec' | 'legacy_cli'; @@ -1323,7 +1321,7 @@ const runClaim = async (claim: Claim) => { ...claim.secrets.map((secret) => secret.value), ].filter(Boolean); const redact = createRedactor(secretValues); - let acquiredBoxUser: string | undefined; + let acquiredBoxHandle: BoxHandle | undefined; try { if ((claim.job.runtime ?? 'opencode') !== 'opencode') { throw new Error('Legacy OpenAI direct jobs are no longer supported.'); @@ -1354,12 +1352,13 @@ const runClaim = async (claim: Claim) => { // Start (or reuse) the persistent per-user box that this thread — and the // terminal — exec into. It mounts the home, so the clone below is visible. - const boxName = await acquireUserBox({ + const boxHandle = await acquireUserBox({ username, workdir: homeDir, containerHome, }); - acquiredBoxUser = username; + acquiredBoxHandle = boxHandle; + const boxName = boxHandle.boxName; const repoDir = await cloneRepository({ workdir: checkoutParent, @@ -1381,6 +1380,7 @@ const runClaim = async (claim: Claim) => { containerRepo, repoDir, boxName, + boxHandle, githubToken, redact, }; @@ -1460,7 +1460,7 @@ const runClaim = async (claim: Claim) => { ).catch((stopError: unknown) => { console.error(stopError); }); - if (acquiredBoxUser) releaseUserBox(acquiredBoxUser); + acquiredBoxHandle?.release(); } }; @@ -1883,7 +1883,7 @@ export const openWorkspacePullRequest = async (jobId: string) => { activeWorkspaces.delete(jobId); // The persistent per-user home + ~/Code checkouts survive across sessions; // release the box (reaped once no other thread/terminal holds it). - releaseUserBox(workspace.username); + workspace.boxHandle.release(); return { pullRequestUrl: pullRequest.html_url, pullRequestNumber: pullRequest.number, @@ -1900,7 +1900,7 @@ export const stopWorkspace = async (jobId: string) => { activeWorkspaces.delete(jobId); // The persistent per-user home + ~/Code checkouts survive across sessions; // release the box (reaped once no other thread/terminal holds it). - releaseUserBox(workspace.username); + workspace.boxHandle.release(); return { success: true }; }; diff --git a/apps/agent-worker/tests/unit/user-container.test.ts b/apps/agent-worker/tests/unit/user-container.test.ts new file mode 100644 index 0000000..ed33aff --- /dev/null +++ b/apps/agent-worker/tests/unit/user-container.test.ts @@ -0,0 +1,85 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; + +vi.mock('../../src/env', () => ({ env: { boxIdleMs: 1000 } })); +vi.mock('../../src/runtime/docker', () => ({ + ensureUserContainer: vi.fn((args: { username: string }) => + Promise.resolve(`spoon-box-${args.username}`), + ), + stopWorkspaceContainer: vi.fn(() => Promise.resolve()), + userContainerName: (username: string) => `spoon-box-${username}`, + listWorkspaceContainerNames: vi.fn(() => Promise.resolve([])), +})); + +const load = async () => await import('../../src/user-container'); + +describe('box registry', () => { + beforeEach(() => vi.useFakeTimers()); + afterEach(async () => { + const module = await load(); + module._resetBoxRegistryForTests(); + vi.useRealTimers(); + vi.resetModules(); + vi.clearAllMocks(); + }); + + test('serializes concurrent acquire into a single docker run', async () => { + const module = await load(); + const docker = await import('../../src/runtime/docker'); + const [a, b] = await Promise.all([ + module.acquireUserBox({ + username: 'alice', + workdir: '/w', + containerHome: '/home/alice', + }), + module.acquireUserBox({ + username: 'alice', + workdir: '/w', + containerHome: '/home/alice', + }), + ]); + expect(docker.ensureUserContainer).toHaveBeenCalledTimes(1); + expect(a.boxName).toBe('spoon-box-alice'); + a.release(); + vi.advanceTimersByTime(2000); + expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled(); + b.release(); + vi.advanceTimersByTime(2000); + expect(docker.stopWorkspaceContainer).toHaveBeenCalledTimes(1); + }); + + test('double-release of one handle does not reap a box held elsewhere', async () => { + const module = await load(); + const docker = await import('../../src/runtime/docker'); + const a = await module.acquireUserBox({ + username: 'bob', + workdir: '/w', + containerHome: '/home/bob', + }); + const b = await module.acquireUserBox({ + username: 'bob', + workdir: '/w', + containerHome: '/home/bob', + }); + a.release(); + a.release(); + vi.advanceTimersByTime(2000); + expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled(); + b.release(); + vi.advanceTimersByTime(2000); + expect(docker.stopWorkspaceContainer).toHaveBeenCalledTimes(1); + }); + + test('reconcileExistingBoxes adopts orphans with zero refs and reaps them', async () => { + const module = await load(); + const docker = await import('../../src/runtime/docker'); + ( + docker.listWorkspaceContainerNames as unknown as ReturnType + ).mockResolvedValueOnce(['spoon-box-carol']); + await module.reconcileExistingBoxes(); + expect(module.runningBoxUsernames().has('carol')).toBe(true); + vi.advanceTimersByTime(2000); + expect(docker.stopWorkspaceContainer).toHaveBeenCalledWith( + 'spoon-box-carol', + ); + }); +});