fix(cloud-agent): make stale workspace cleanup runtime-aware (#3656)

Author: eshurakov

services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.test.ts

@@ -146,6 +146,78 @@ describe('CloudflareAgentSandbox', () => {
     ensureBootstrapWrapper.mockRestore();
   });
 
+  it('reclaims stale bootstrap workspaces without inspecting Docker', async () => {
+    const bootstrapSession = {};
+    const createSession = vi.fn().mockResolvedValue(bootstrapSession);
+    const ensureBootstrapWrapper = vi
+      .spyOn(WrapperClient, 'ensureBootstrapWrapper')
+      .mockResolvedValueOnce({ client: {} as WrapperClient });
+    const exec = vi
+      .fn()
+      .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '536870912  10485760000\n', stderr: '' })
+      .mockResolvedValueOnce({
+        exitCode: 0,
+        stdout: 'agent_stale-aaaa\nagent_cloudflare\n',
+        stderr: '',
+      })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '0\n', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '3145728000  10485760000\n', stderr: '' });
+    const sandbox = new CloudflareAgentSandbox({} as Env, metadata(), {
+      resolveSandbox: () =>
+        ({
+          exec,
+          listProcesses: vi.fn().mockResolvedValue([]),
+          createSession,
+        }) as unknown as SandboxInstance,
+    });
+
+    await expect(sandbox.ensureWrapper(ensureRequest())).resolves.toMatchObject({
+      status: 'wrapper-running',
+    });
+    expect(exec.mock.calls.every(call => !call[0].includes('docker'))).toBe(true);
+    expect(createSession).toHaveBeenCalled();
+    ensureBootstrapWrapper.mockRestore();
+  });
+
+  it('keeps unresolved DIND bootstrap cleanup fail-closed', async () => {
+    const unresolvedDindMetadata = {
+      ...metadata(),
+      workspace: { sandboxId: 'dind-unresolved' },
+    } satisfies SessionMetadata;
+    const request = ensureRequest();
+    request.plan.workspace = { sandboxId: 'dind-unresolved', metadata: unresolvedDindMetadata };
+    const exec = vi
+      .fn()
+      .mockResolvedValueOnce({ exitCode: 1, stdout: '', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '536870912  10485760000\n', stderr: '' })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: 'agent_stale-aaaa\n', stderr: '' })
+      .mockResolvedValueOnce({
+        exitCode: 0,
+        stdout: '/run/user/1000/docker.sock',
+        stderr: '',
+      })
+      .mockRejectedValueOnce(new Error('docker inspection unavailable'))
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '536870912  10485760000\n', stderr: '' });
+    const sandbox = new CloudflareAgentSandbox({} as Env, unresolvedDindMetadata, {
+      resolveSandbox: () =>
+        ({
+          exec,
+          listProcesses: vi.fn().mockResolvedValue([]),
+          createSession: vi.fn(),
+        }) as unknown as SandboxInstance,
+    });
+
+    await expect(sandbox.ensureWrapper(request)).rejects.toBeInstanceOf(
+      WorkspaceCapacityAdmissionRejectedError
+    );
+    expect(exec.mock.calls[4][0]).toContain('docker ps');
+    expect(exec.mock.calls.every(call => !call[0].includes('stat'))).toBe(true);
+    expect(exec.mock.calls.every(call => !call[0].includes('rm -rf'))).toBe(true);
+  });
+
   it('passes a leased physical identity into bootstrap startup', async () => {
     const bootstrapSession = {};
     const ensureBootstrapWrapper = vi

services/cloud-agent-next/src/agent-sandbox/cloudflare/cloudflare-agent-sandbox.ts

@@ -230,7 +230,9 @@ export class CloudflareAgentSandbox implements AgentSandbox {
     const workspaceWarm = await this.workspaceHasGit(sandbox, prepared.context.workspacePath);
     if (!workspaceWarm) {
       request.onProgress?.('disk_check', 'Checking disk space...');
-      await checkDiskAndCleanBeforeSetup(sandbox, orgId, userId, sessionId);
+      await checkDiskAndCleanBeforeSetup(sandbox, orgId, userId, sessionId, {
+        inspectContainers: sandboxId.startsWith('dind-'),
+      });
     }
     request.onProgress?.('kilo_server', 'Starting Kilo...');
     const bootstrapSession = await sandbox.createSession({

services/cloud-agent-next/src/kilo/wrapper-manager.ts

@@ -164,7 +164,7 @@ export async function findWrapperForSession(
  * `kilo.agentSession=<id>`. The published port we want is buried in the
  * `Ports` column (`0.0.0.0:5xxx->5xxx/tcp` or `127.0.0.1:5xxx->5xxx/tcp`).
  */
-type LabeledWrapperRow = {
+export type LabeledWrapperRow = {
   containerId: string;
   agentSessionId: string;
   port?: number;

services/cloud-agent-next/src/session-service.test.ts

@@ -286,6 +286,13 @@ describe('SessionService.prepareWorkspace', () => {
       onProgress: progress,
     });
 
+    expect(workspaceMocks.checkDiskAndCleanBeforeSetup).toHaveBeenCalledWith(
+      sandbox,
+      undefined,
+      'user_test',
+      'agent_test',
+      { inspectContainers: false }
+    );
     expect(workspaceMocks.cloneGitRepo).toHaveBeenCalledWith(
       session,
       '/workspace/user/sessions/agent_test',
@@ -375,11 +382,57 @@ describe('SessionService.prepareWorkspace', () => {
       })
     ).rejects.toBe(rejection);
 
+    expect(workspaceMocks.checkDiskAndCleanBeforeSetup).toHaveBeenCalledWith(
+      sandbox,
+      undefined,
+      'user_test',
+      'agent_test',
+      { inspectContainers: true }
+    );
     expect(workspaceMocks.setupWorkspace).not.toHaveBeenCalled();
     expect(sandbox.createSessionMock).not.toHaveBeenCalled();
     expect(devcontainerMocks.bringUpDevContainer).not.toHaveBeenCalled();
   });
 
+  it('keeps requested devcontainer cleanup fail-closed when the sandbox ID is not DIND', async () => {
+    const session = createSession(false);
+    const sandbox = createSandbox(session);
+    const metadata = {
+      ...createMetadata(),
+      workspace: {
+        sandboxId: 'usr-abcdef' as const,
+        devcontainerRequested: true,
+      },
+    } satisfies CloudAgentSessionState;
+    const rejection = new WorkspaceCapacityAdmissionRejectedError({
+      availableMB: 512,
+      thresholdMB: 2048,
+      cleaned: 0,
+      skipped: 1,
+    });
+    workspaceMocks.checkDiskAndCleanBeforeSetup.mockRejectedValueOnce(rejection);
+
+    await expect(
+      new SessionService().prepareWorkspace({
+        sandbox,
+        sandboxId: 'usr-abcdef',
+        userId: 'user_test',
+        sessionId: 'agent_test' as SessionId,
+        env: createEnv(),
+        metadata,
+        kilocodeModel: 'test-model',
+      })
+    ).rejects.toBe(rejection);
+
+    expect(workspaceMocks.checkDiskAndCleanBeforeSetup).toHaveBeenCalledWith(
+      sandbox,
+      undefined,
+      'user_test',
+      'agent_test',
+      { inspectContainers: true }
+    );
+  });
+
   it('hydrates requested devcontainer metadata while preparing a cold DIND workspace', async () => {
     const session = createSession(false);
     const writeFile = vi.fn().mockResolvedValue(undefined);

services/cloud-agent-next/src/session-service.ts

@@ -1703,7 +1703,12 @@ export class SessionService {
     }
 
     onProgress?.('disk_check', 'Checking disk space…');
-    await checkDiskAndCleanBeforeSetup(sandbox, orgId, userId, sessionId);
+    await checkDiskAndCleanBeforeSetup(sandbox, orgId, userId, sessionId, {
+      inspectContainers:
+        sandboxId.startsWith('dind-') ||
+        metadata.workspace?.devcontainerRequested === true ||
+        metadata.devcontainer !== undefined,
+    });
 
     onProgress?.('workspace_setup', 'Setting up workspace…');
     await setupWorkspace(sandbox, userId, orgId, sessionId);