spoon/apps/agent-worker/tests/unit/user-container.test.ts

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('passes an onCreated hook so a fresh box gets its user initialized', async () => {
    const module = await load();
    const docker = await import('../../src/runtime/docker');
    await module.acquireUserBox({
      username: 'ada',
      workdir: '/w',
      containerHome: '/home/ada',
    });
    const call = vi.mocked(docker.ensureUserContainer).mock.calls[0]?.[0] as {
      onCreated?: unknown;
    };
    expect(call.onCreated).toBeTypeOf('function');
  });

  test('re-ensures the container on every acquire so an externally removed box self-heals', async () => {
    const module = await load();
    const docker = await import('../../src/runtime/docker');
    const first = await module.acquireUserBox({
      username: 'zoe',
      workdir: '/w',
      containerHome: '/home/zoe',
    });
    first.release();
    // Simulate `docker rm` behind the worker's back: the registry entry still
    // exists, but acquire must go back to the runtime rather than trust it.
    const second = await module.acquireUserBox({
      username: 'zoe',
      workdir: '/w',
      containerHome: '/home/zoe',
    });
    expect(docker.ensureUserContainer).toHaveBeenCalledTimes(2);
    expect(second.boxName).toBe('spoon-box-zoe');
  });

  test('serializes concurrent acquires so ensure calls never overlap', async () => {
    const module = await load();
    const docker = await import('../../src/runtime/docker');
    let inFlight = 0;
    vi.mocked(docker.ensureUserContainer).mockImplementation(
      async (args: { username: string }) => {
        inFlight += 1;
        expect(inFlight).toBe(1);
        await Promise.resolve();
        inFlight -= 1;
        return `spoon-box-${args.username}`;
      },
    );
    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(a.boxName).toBe('spoon-box-alice');
    a.release();
    await vi.advanceTimersByTimeAsync(2000);
    expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled();
    b.release();
    await vi.advanceTimersByTimeAsync(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();
    await vi.advanceTimersByTimeAsync(2000);
    expect(docker.stopWorkspaceContainer).not.toHaveBeenCalled();
    b.release();
    await vi.advanceTimersByTimeAsync(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);
    await vi.advanceTimersByTimeAsync(2000);
    expect(docker.stopWorkspaceContainer).toHaveBeenCalledWith(
      'spoon-box-carol',
    );
  });

  test('acquiring an adopted box ensures it and returns valid handles', async () => {
    const module = await load();
    const docker = await import('../../src/runtime/docker');
    vi.mocked(docker.listWorkspaceContainerNames).mockResolvedValueOnce([
      'spoon-box-casey',
    ]);
    await module.reconcileExistingBoxes();

    const first = await module.acquireUserBox({
      username: 'casey',
      workdir: '/w',
      containerHome: '/home/casey',
    });
    const second = await module.acquireUserBox({
      username: 'casey',
      workdir: '/w',
      containerHome: '/home/casey',
    });

    expect(docker.ensureUserContainer).toHaveBeenCalled();
    expect(first.boxName).toBe('spoon-box-casey');
    expect(second.boxName).toBe('spoon-box-casey');
  });

  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']);
  });
});