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

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