fix(worker): make user-box lifecycle concurrency-safe

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();

apps/agent-worker/src/runtime/docker.ts

@@ -362,6 +362,7 @@ export const ensureUserContainer = async (args: {
     [
       'run',
       '-d',
+      '--init',
       '--name',
       name,
       '--memory',

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.

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<string, Box>();
+// 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<string> => {
-  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<symbol>;
+  idleTimer?: ReturnType<typeof setTimeout>;
+};
+const boxes = new Map<string, Box>();
+const locks = new Map<string, Promise<unknown>>();
+
+// Per-username async mutex: chain each operation after the previous one.
+const withLock = <T>(username: string, fn: () => Promise<T>): Promise<T> => {
+  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<BoxHandle> =>
+  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<void> => {
+  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<string> => new Set(boxes.keys());
+
+export const _resetBoxRegistryForTests = () => {
+  for (const box of boxes.values()) {
+    if (box.idleTimer) clearTimeout(box.idleTimer);
+  }
+  boxes.clear();
+  locks.clear();
+};

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 };
 };
 

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<typeof vi.fn>
+    ).mockResolvedValueOnce(['spoon-box-carol']);
+    await module.reconcileExistingBoxes();
+    expect(module.runningBoxUsernames().has('carol')).toBe(true);
+    vi.advanceTimersByTime(2000);
+    expect(docker.stopWorkspaceContainer).toHaveBeenCalledWith(
+      'spoon-box-carol',
+    );
+  });
+});