From 284f7ea52b043f9ecb372d3d6bdbe25ec8a9add4 Mon Sep 17 00:00:00 2001
From: Gabriel Brown <gbrown@ksensetech.com>
Date: Fri, 10 Jul 2026 17:00:22 -0400
Subject: [PATCH] fix(worker): serialize box failure and idle reaping

---
 apps/agent-worker/src/user-container.ts       | 14 +++-
 .../tests/unit/user-container.test.ts         | 79 +++++++++++++++++--
 2 files changed, 86 insertions(+), 7 deletions(-)

diff --git a/apps/agent-worker/src/user-container.ts b/apps/agent-worker/src/user-container.ts
index 4f119a2..3b6890f 100644
--- a/apps/agent-worker/src/user-container.ts
+++ b/apps/agent-worker/src/user-container.ts
@@ -38,8 +38,14 @@ const scheduleReapIfIdle = (username: string) => {
   if (!box || box.refs.size > 0) return;
   if (box.idleTimer) clearTimeout(box.idleTimer);
   box.idleTimer = setTimeout(() => {
-    void stopWorkspaceContainer(userContainerName(username));
-    boxes.delete(username);
+    void withLock(username, async () => {
+      const current = boxes.get(username);
+      if (current !== box || current.refs.size > 0) return;
+      await stopWorkspaceContainer(current.name);
+      if (boxes.get(username) === current && current.refs.size === 0) {
+        boxes.delete(username);
+      }
+    });
   }, env.boxIdleMs);
 };
 
@@ -83,6 +89,10 @@ export const acquireUserBox = (args: {
       return handle;
     } catch (error) {
       handle.release();
+      if (boxes.get(args.username) === box && box.refs.size === 0) {
+        if (box.idleTimer) clearTimeout(box.idleTimer);
+        boxes.delete(args.username);
+      }
       throw error;
     }
   });
diff --git a/apps/agent-worker/tests/unit/user-container.test.ts b/apps/agent-worker/tests/unit/user-container.test.ts
index ed33aff..08b36ad 100644
--- a/apps/agent-worker/tests/unit/user-container.test.ts
+++ b/apps/agent-worker/tests/unit/user-container.test.ts
@@ -40,10 +40,10 @@ describe('box registry', () => {
     expect(docker.ensureUserContainer).toHaveBeenCalledTimes(1);
     expect(a.boxName).toBe('spoon-box-alice');
     a.release();
-    vi.advanceTimersByTime(2000);
+    await vi.advanceTimersByTimeAsync(2000);
     expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled();
     b.release();
-    vi.advanceTimersByTime(2000);
+    await vi.advanceTimersByTimeAsync(2000);
     expect(docker.stopWorkspaceContainer).toHaveBeenCalledTimes(1);
   });
 
@@ -62,10 +62,10 @@ describe('box registry', () => {
     });
     a.release();
     a.release();
-    vi.advanceTimersByTime(2000);
+    await vi.advanceTimersByTimeAsync(2000);
     expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled();
     b.release();
-    vi.advanceTimersByTime(2000);
+    await vi.advanceTimersByTimeAsync(2000);
     expect(docker.stopWorkspaceContainer).toHaveBeenCalledTimes(1);
   });
 
@@ -77,9 +77,78 @@ describe('box registry', () => {
     ).mockResolvedValueOnce(['spoon-box-carol']);
     await module.reconcileExistingBoxes();
     expect(module.runningBoxUsernames().has('carol')).toBe(true);
-    vi.advanceTimersByTime(2000);
+    await vi.advanceTimersByTimeAsync(2000);
     expect(docker.stopWorkspaceContainer).toHaveBeenCalledWith(
       'spoon-box-carol',
     );
   });
+
+  test('retries ensure after the first acquire fails', async () => {
+    const module = await load();
+    const docker = await import('../../src/runtime/docker');
+    const ensure = vi.mocked(docker.ensureUserContainer);
+    ensure
+      .mockRejectedValueOnce(new Error('docker failed'))
+      .mockResolvedValueOnce('spoon-box-dana');
+
+    await expect(
+      module.acquireUserBox({
+        username: 'dana',
+        workdir: '/w',
+        containerHome: '/home/dana',
+      }),
+    ).rejects.toThrow('docker failed');
+
+    const handle = await module.acquireUserBox({
+      username: 'dana',
+      workdir: '/w',
+      containerHome: '/home/dana',
+    });
+    expect(ensure).toHaveBeenCalledTimes(2);
+    expect(handle.boxName).toBe('spoon-box-dana');
+  });
+
+  test('finishes an old idle reap before a new acquire can ensure the box', async () => {
+    const module = await load();
+    const docker = await import('../../src/runtime/docker');
+    const ensure = vi.mocked(docker.ensureUserContainer);
+    const stop = vi.mocked(docker.stopWorkspaceContainer);
+    const events: string[] = [];
+    let finishStop!: () => void;
+    const stopDone = new Promise<void>((resolve) => {
+      finishStop = resolve;
+    });
+    ensure.mockReset().mockImplementation((args: { username: string }) => {
+      events.push('ensure');
+      return Promise.resolve(`spoon-box-${args.username}`);
+    });
+    stop.mockReset().mockImplementationOnce(() => {
+      events.push('stop:start');
+      return stopDone.then(() => {
+        events.push('stop:end');
+      });
+    });
+
+    const first = await module.acquireUserBox({
+      username: 'erin',
+      workdir: '/w',
+      containerHome: '/home/erin',
+    });
+    first.release();
+    await vi.advanceTimersByTimeAsync(1000);
+
+    const secondPromise = module.acquireUserBox({
+      username: 'erin',
+      workdir: '/w',
+      containerHome: '/home/erin',
+    });
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(events).toEqual(['ensure', 'stop:start']);
+
+    finishStop();
+    const second = await secondPromise;
+    expect(second.boxName).toBe('spoon-box-erin');
+    expect(events).toEqual(['ensure', 'stop:start', 'stop:end', 'ensure']);
+  });
 });