@fitlab-ai/agent-infra 0.9.5 → 0.9.6

package/lib/sandbox/commands/create.ts

@@ -2,10 +2,9 @@ import fs from 'node:fs';
 import os from 'node:os';
 import path from 'node:path';
 import { createHash } from 'node:crypto';
-import { execFileSync, spawn } from 'node:child_process';
+import { execFileSync } from 'node:child_process';
 import type { ExecFileSyncOptions, StdioOptions } from 'node:child_process';
 import { parseArgs } from 'node:util';
-import { fileURLToPath } from 'node:url';
 import * as p from '@clack/prompts';
 import pc from 'picocolors';
 import { listAgentClientAdapters } from '../../agent-clients/registry.ts';
@@ -67,14 +66,16 @@ import type { SandboxTool, TmpfsSeedEntry } from '../tools.ts';
 import {
   assertFreshSandboxReady,
   hydrateTmpfsSeedEntries,
-  prepareTmpfsMounts
+  prepareTmpfsMounts,
+  startSandboxControlBroker
 } from '../recovery.ts';
 import { hostJoin, toEnginePath, volumeArg } from '../engines/wsl2-paths.ts';
 import { sandboxCoreBindMounts } from '../mounts.ts';
 import {
   assertSandboxTaskSource,
   materializeSandboxControl,
-  materializeSandboxWorkspaceView
+  materializeSandboxWorkspaceView,
+  prepareSandboxWorkspaceMountTargets
 } from '../workspace-view.ts';
 import { clipboardHostDir, CONTAINER_CLIPBOARD_MOUNT } from '../clipboard/paths.ts';
 import { validateSelinuxDisableEnv } from '../engines/selinux.ts';
@@ -157,12 +158,6 @@ type HostShellConfig = {
   mounts: Array<{ hostPath: string; containerPath: string }>;
 };
 
-function internalCliPath(): string {
-  const directory = path.dirname(fileURLToPath(import.meta.url));
-  const extension = path.extname(fileURLToPath(import.meta.url));
-  return path.resolve(directory, '..', '..', '..', 'bin', `internal-cli${extension}`);
-}
-
 function resolveToolDirs(config: Pick<SandboxCreateConfig, 'project'>, tools: SandboxTool[], branch: string): ResolvedTool[] {
   return tools.map((tool) => {
     const candidates = toolConfigDirCandidates(tool, config.project, branch);
@@ -1254,9 +1249,13 @@ export async function create(args: string[]): Promise<void> {
                 const requestedDescription = target.workspace.mode === 'task-bound'
                   ? `task-bound:${target.workspace.taskId}`
                   : target.workspace.mode;
+                const recoveryTarget = target.workspace.mode === 'task-bound'
+                  ? target.workspace.taskId
+                  : branch;
                 throw new Error(
                   `SANDBOX_WORKSPACE_IDENTITY_CONFLICT: container '${name}' is ${existingDescription}, `
-                  + `but this request is ${requestedDescription}. Run 'ai sandbox rm ${branch}' and retry.`
+                  + `but this request is ${requestedDescription}. `
+                  + `Run 'ai sandbox start --recreate ${recoveryTarget}' on the host, then retry.`
                 );
               }
             }
@@ -1342,6 +1341,7 @@ export async function create(args: string[]): Promise<void> {
               container,
               identity: target.workspace
             });
+            prepareSandboxWorkspaceMountTargets(worktree);
             const control = materializeSandboxControl({
               base: effectiveConfig.controlBase,
               repoRoot: effectiveConfig.repoRoot,
@@ -1362,6 +1362,7 @@ export async function create(args: string[]): Promise<void> {
               shellConfigHostDir: hostShellConfig.hostDir,
               workspaceViewRoot: workspaceView.root,
               controlDir: control.channelDir,
+              controlStatusDir: control.statusDir,
               ...(target.workspace.mode === 'task-bound'
                 ? {
                   taskSource: assertSandboxTaskSource(
@@ -1445,7 +1446,11 @@ export async function create(args: string[]): Promise<void> {
               '-e',
               `AGENT_INFRA_CONTROL_TOKEN=${control.token}`,
               '-e',
+              `AGENT_INFRA_CONTROL_GENERATION=${control.generation}`,
+              '-e',
               'AGENT_INFRA_CONTROL_DIR=/run/agent-infra/control',
+              '-e',
+              'AGENT_INFRA_CONTROL_STATUS_DIR=/run/agent-infra/control-status',
               '--label',
               `${sandboxRuntimeCapabilityLabel(effectiveConfig)}=${capabilityPlan.runtimeSignature}`,
               ...coreVolumes,
@@ -1467,16 +1472,7 @@ export async function create(args: string[]): Promise<void> {
               '/workspace',
               effectiveConfig.imageName
             ]);
-            const broker = spawn(
-              process.execPath,
-              [internalCliPath(), 'sandbox-control', 'serve', '--manifest', control.manifestPath],
-              { cwd: effectiveConfig.repoRoot, detached: true, stdio: 'ignore' }
-            );
-            broker.once('error', () => {
-              // Fresh readiness verification below remains the source of truth;
-              // consume detached spawn errors so they cannot crash the command.
-            });
-            broker.unref();
+            await startSandboxControlBroker(effectiveConfig.repoRoot, control.manifestPath);
             createdTmpfsSeedPlan = tmpfsSeedPlan;
           } finally {
             envFile.cleanup();

package/lib/sandbox/commands/prune.ts

@@ -6,16 +6,27 @@ import pc from 'picocolors';
 import { loadConfig } from '../config.ts';
 import type { SandboxConfig } from '../config.ts';
 import { safeNameCandidates, sandboxBranchLabel, sandboxLabel } from '../constants.ts';
-import { detectEngine } from '../engine.ts';
+import { detectEngine, ENGINES } from '../engine.ts';
 import { hostJoin } from '../engines/wsl2-paths.ts';
 import { removeManagedDir, removeWorktreeDir } from '../managed-fs.ts';
 import { parseLabels } from './ls.ts';
 import { runEngine, runSafe } from '../shell.ts';
 import { createSandboxCapabilityPlan } from '../agent-client-reconciler.ts';
 import type { SandboxTool } from '../tools.ts';
+import { createCleanPermit, formatWorktreeSnapshot, inspectWorktrees } from '../worktree-safety.ts';
+import type { WorktreeRemovalPermit } from '../worktree-safety.ts';
 
 const USAGE = `Usage: ai sandbox prune [--dry-run]`;
 
+type PruneDependencies = {
+  config?: SandboxConfig;
+  tools?: SandboxTool[];
+  engine?: string;
+  labelsOutput?: string;
+  confirm?: typeof p.confirm;
+  isCancel?: typeof p.isCancel;
+};
+
 type OrphanKind = 'shell' | 'worktree' | 'share' | 'tool';
 
 export type OrphanGroup = {
@@ -107,12 +118,19 @@ export function collectOrphanGroups(
   return groups;
 }
 
-export function removeOrphanGroups(config: SandboxConfig, groups: OrphanGroup[]): boolean {
+export function removeOrphanGroups(
+  config: SandboxConfig,
+  groups: OrphanGroup[],
+  permits: ReadonlyMap<string, WorktreeRemovalPermit>,
+  allowRegisteredPathFallback = false
+): boolean {
   let removedWorktrees = false;
   for (const group of groups) {
     for (const dir of group.dirs) {
       if (group.kind === 'worktree') {
-        removeWorktreeDir(config.repoRoot, group.base, dir);
+        const permit = permits.get(path.resolve(dir));
+        if (!permit) throw new Error(`Missing worktree removal permit: ${dir}`);
+        removeWorktreeDir(config.repoRoot, group.base, dir, permit, { allowRegisteredPathFallback });
         removedWorktrees = true;
       } else {
         removeManagedDir(group.base, dir);
@@ -142,7 +160,7 @@ function writeGroups(groups: OrphanGroup[]): void {
   }
 }
 
-export async function prune(args: string[]): Promise<void> {
+export async function prune(args: string[], dependencies: PruneDependencies = {}): Promise<void> {
   const { values } = parseArgs({
     args,
     strict: true,
@@ -157,9 +175,9 @@ export async function prune(args: string[]): Promise<void> {
     return;
   }
 
-  const config = loadConfig();
-  const tools = [...createSandboxCapabilityPlan(config).cleanupInventory];
-  const engine = detectEngine(config);
+  const config = dependencies.config ?? loadConfig();
+  const tools = dependencies.tools ?? [...createSandboxCapabilityPlan(config).cleanupInventory];
+  const engine = dependencies.engine ?? detectEngine(config);
   const psArgs = [
     'ps',
     '-a',
@@ -168,14 +186,25 @@ export async function prune(args: string[]): Promise<void> {
     '--format',
     '{{.Labels}}'
   ];
-  let labelsOutput: string;
-  try {
-    labelsOutput = runEngine(engine, 'docker', psArgs);
-  } catch {
-    throw new Error('Unable to determine active sandbox branches: docker ps failed');
+  let labelsOutput = dependencies.labelsOutput;
+  if (labelsOutput === undefined) {
+    try {
+      labelsOutput = runEngine(engine, 'docker', psArgs);
+    } catch {
+      throw new Error('Unable to determine active sandbox branches: docker ps failed');
+    }
   }
+  const confirm = dependencies.confirm ?? p.confirm;
+  const isCancel = dependencies.isCancel ?? p.isCancel;
   const groups = collectOrphanGroups(config, tools, activeBranchesFromLabels(config, labelsOutput));
   const count = orphanCount(groups);
+  const worktrees = groups.filter((group) => group.kind === 'worktree').flatMap((group) => group.dirs);
+  const inspections = inspectWorktrees(worktrees);
+  const blockers = inspections.filter((inspection) => inspection.status !== 'clean');
+  const permits = new Map<string, WorktreeRemovalPermit>();
+  for (const inspection of inspections) {
+    if (inspection.status === 'clean') permits.set(inspection.snapshot.worktree, createCleanPermit(inspection.snapshot));
+  }
 
   p.intro(pc.cyan(`Pruning orphaned sandbox state for ${config.project}`));
 
@@ -186,23 +215,32 @@ export async function prune(args: string[]): Promise<void> {
   }
 
   writeGroups(groups);
+  for (const blocker of blockers) {
+    p.log.error(blocker.status === 'failed'
+      ? `${JSON.stringify(blocker.worktree)}: ${blocker.message}`
+      : formatWorktreeSnapshot(blocker.snapshot));
+  }
 
   if (values['dry-run']) {
     p.outro(pc.green('Dry run complete'));
     return;
   }
 
-  const shouldRemove = await p.confirm({
+  if (blockers.length > 0) {
+    throw new Error('Refusing to prune because worktree preflight found blocker(s); no orphan state was deleted');
+  }
+
+  const shouldRemove = await confirm({
     message: `Remove ${count} orphaned sandbox state dirs?`,
     initialValue: true
   });
 
-  if (p.isCancel(shouldRemove) || !shouldRemove) {
+  if (isCancel(shouldRemove) || !shouldRemove) {
     p.outro('Cancelled');
     return;
   }
 
-  const removedWorktrees = removeOrphanGroups(config, groups);
+  const removedWorktrees = removeOrphanGroups(config, groups, permits, engine === ENGINES.WSL2);
   if (removedWorktrees) {
     runSafe('git', ['-C', config.repoRoot, 'worktree', 'prune']);
   }