@fitlab-ai/agent-infra 0.9.5 → 0.9.7

package/lib/sandbox/recovery.ts

@@ -2,13 +2,15 @@ import fs from 'node:fs';
 import path from 'node:path';
 import { spawn } from 'node:child_process';
 import { fileURLToPath } from 'node:url';
+import { nodeEntryArgs } from './control/executor.ts';
 import {
   containerNameCandidates,
   sandboxBranchLabel,
   sandboxLabel,
   sandboxRuntimeCapabilityLabel,
   sandboxTaskIdLabel,
-  sandboxWorkspaceModeLabel
+  sandboxWorkspaceModeLabel,
+  worktreeDirCandidates
 } from './constants.ts';
 import { isAgentClientId } from '../agent-clients/types.ts';
 import {
@@ -29,9 +31,11 @@ import {
 import {
   parseSandboxWorkspaceIdentity,
   sameSandboxWorkspaceIdentity,
+  type SandboxContainerWorkspaceIdentity,
   type SandboxWorkspaceIdentity
 } from './workspace-identity.ts';
-import { runEngine, runOkEngine, runVerboseEngine } from './shell.ts';
+import { inspectWorktree, type WorktreeSnapshot } from './worktree-safety.ts';
+import { commandForEngine, runEngine, runOkEngine, runProbe, runVerboseEngine } from './shell.ts';
 import {
   declaredTmpfsSeedEntries,
   resolveTools,
@@ -45,10 +49,28 @@ import {
   startSandboxContainer,
   type SandboxRow
 } from './commands/list-running.ts';
-import { processIdentityMatches } from '../server/process-state.ts';
+import { getProcessStartTime, processIdentityMatches, removePidFileIfMatches } from '../server/process-state.ts';
+import {
+  acquireSandboxControlBrokerStartup,
+  garbageCollectSandboxControlRoot,
+  isSandboxControlRootQuiescing,
+  readSandboxControlManifest,
+  type BrokerOwner
+} from './control/lifecycle.ts';
+import {
+  appendSandboxControlAudit,
+  readSandboxControlStatus
+} from './control/state.ts';
+import { inspectSandboxControlContainer } from './control/container-identity.ts';
+import {
+  SANDBOX_CONTROL_FUTURE_SKEW_MS,
+  SANDBOX_CONTROL_STATUS_STALE_MS,
+  type SandboxControlManifest
+} from './control/protocol.ts';
 
 export type SandboxRecoveryFinding = {
   repairKind: 'permissions' | 'missing-seed' | 'builtin-link' | 'hard-failure';
+  code?: string;
   message: string;
   path?: string;
   seed?: TmpfsSeedEntry;
@@ -66,6 +88,11 @@ type SandboxAgentClientCheckSnapshot = {
 
 export type SandboxRecoverySnapshot = {
   identityOk: boolean;
+  containerIdValid: boolean;
+  expectedBranch: string;
+  actualBranch: string | null;
+  expectedWorkspace: SandboxWorkspaceIdentity;
+  actualWorkspace: SandboxContainerWorkspaceIdentity;
   runtimeCapabilityOk?: boolean;
   unexpectedCapabilityMounts?: string[];
   mounts: Array<{
@@ -118,8 +145,28 @@ type RecoveryCommandDeps = {
   runVerbose?: typeof runVerboseEngine;
   start?: typeof startSandboxContainer;
   fetchRows?: typeof fetchSandboxRows;
+  ensureControlBroker?: () => void | Promise<void>;
+  probe?: typeof runProbe;
 };
 
+export function worktreeProbeForEngine(
+  engine: string,
+  probe: typeof runProbe = runProbe
+): typeof runProbe {
+  return (cmd, args, opts = {}) => {
+    const engineArgs = [...args];
+    if (engine === 'wsl2' && cmd === 'git') {
+      for (let index = 0; index < engineArgs.length - 1; index += 1) {
+        if (engineArgs[index] === '-C') {
+          engineArgs[index + 1] = toEnginePath(engine, engineArgs[index + 1]!);
+        }
+      }
+    }
+    const command = commandForEngine(engine, cmd, engineArgs);
+    return probe(command.cmd, command.args, opts);
+  };
+}
+
 type EnsureSandboxReadyParams = {
   config: SandboxConfig;
   engine: string;
@@ -127,6 +174,7 @@ type EnsureSandboxReadyParams = {
   workspace?: SandboxWorkspaceIdentity;
   row: SandboxRow;
   allowRecreate?: boolean;
+  forceRecreate?: boolean;
   recreate?: (branch: string) => Promise<void>;
   writeWarning?: (message: string) => void;
   deps?: RecoveryCommandDeps;
@@ -139,11 +187,140 @@ type ExpectedMount = {
   expectedRW: boolean;
 };
 
-function ensureSandboxControlBroker(params: {
+async function waitForSandboxControlBrokerStatus(params: {
+  root: string;
+  statusDir: string;
+  generation: string;
+  pid: number;
+  startTime?: number;
+}, timeoutMs = 5_000): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    if (isSandboxControlRootQuiescing(params.root)) throw new Error('SANDBOX_CONTROL_QUIESCING');
+    try {
+      const status = readSandboxControlStatus(params.statusDir);
+      if (status.generation === params.generation && status.broker.pid === params.pid
+        && (!params.startTime || status.broker.startTime === params.startTime)
+        && Date.now() - status.updatedAt <= SANDBOX_CONTROL_STATUS_STALE_MS
+        && status.updatedAt <= Date.now() + SANDBOX_CONTROL_FUTURE_SKEW_MS) return;
+    } catch {
+      // The broker may still be publishing its initial atomic status.
+    }
+    await new Promise((resolve) => setTimeout(resolve, 25));
+  }
+  throw new Error('SANDBOX_CONTROL_BROKER_START_TIMEOUT');
+}
+
+export async function startSandboxControlBroker(repoRoot: string, manifestPath: string): Promise<void> {
+  const directory = path.dirname(fileURLToPath(import.meta.url));
+  const extension = path.extname(fileURLToPath(import.meta.url));
+  const internalCli = path.resolve(directory, '..', '..', 'bin', `internal-cli${extension}`);
+  const manifest = readSandboxControlManifest(manifestPath);
+  const root = path.dirname(manifestPath);
+  if (isSandboxControlRootQuiescing(root)) throw new Error('SANDBOX_CONTROL_QUIESCING');
+  const brokerPath = path.join(root, 'broker.json');
+  let brokerSnapshot: string | null = null;
+  let replacedBrokerRecord = false;
+  try {
+    brokerSnapshot = fs.readFileSync(brokerPath, 'utf8');
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error;
+  }
+  if (brokerSnapshot !== null) {
+    type ExistingBrokerRecord = {
+      version?: unknown;
+      pid?: unknown;
+      startTime?: unknown;
+      brokerId?: unknown;
+      token?: unknown;
+      generation?: unknown;
+    };
+    let broker: ExistingBrokerRecord | null = null;
+    try {
+      broker = JSON.parse(brokerSnapshot) as ExistingBrokerRecord;
+    } catch {
+      // Malformed owner records are stale and replaced below.
+    }
+    const owner = broker !== null && broker.version === 3 && typeof broker.pid === 'number'
+      && Number.isSafeInteger(broker.pid) && broker.pid > 0 && typeof broker.startTime === 'number'
+      && Number.isSafeInteger(broker.startTime) && typeof broker.brokerId === 'string'
+      && broker.brokerId.length > 0 && typeof broker.token === 'string'
+      && typeof broker.generation === 'string'
+      ? broker as BrokerOwner
+      : null;
+    if (broker === null || broker.version !== 3 || owner === null) {
+      throw new Error('SANDBOX_CONTROL_BROKER_VERSION_INVALID: rebuild the sandbox container before starting a new broker');
+    }
+    const live = owner !== null && processIdentityMatches({ pid: owner.pid, startTime: owner.startTime });
+    if (
+      live
+      && owner !== null
+      && broker !== null
+      && broker.version === 3
+      && broker.token === manifest.token
+      && broker.generation === manifest.generation
+    ) {
+      await waitForSandboxControlBrokerStatus({
+        root,
+        statusDir: manifest.publicStatusDir,
+        generation: manifest.generation,
+        pid: owner.pid,
+        startTime: owner.startTime
+      });
+      return;
+    }
+  }
+  const startupStartTime = getProcessStartTime(process.pid);
+  if (!startupStartTime) throw new Error('SANDBOX_CONTROL_BROKER_IDENTITY_UNAVAILABLE');
+  const releaseStartup = await acquireSandboxControlBrokerStartup(root, {
+    pid: process.pid,
+    startTime: startupStartTime
+  });
+  let child: ReturnType<typeof spawn>;
+  try {
+    const currentBroker = fs.existsSync(brokerPath) ? fs.readFileSync(brokerPath, 'utf8') : null;
+    if (currentBroker !== brokerSnapshot) return startSandboxControlBroker(repoRoot, manifestPath);
+    if (brokerSnapshot !== null) {
+      replacedBrokerRecord = removePidFileIfMatches(brokerPath, brokerSnapshot);
+      if (!replacedBrokerRecord && fs.existsSync(brokerPath)) {
+        throw new Error('SANDBOX_CONTROL_BROKER_OWNER_TRANSITION');
+      }
+    }
+    if (replacedBrokerRecord) {
+      appendSandboxControlAudit(manifest, 'broker-observed-crash');
+      appendSandboxControlAudit(manifest, 'broker-restart');
+    }
+    child = spawn(
+      process.execPath,
+      nodeEntryArgs(internalCli, ['sandbox-control', 'serve', '--manifest', manifestPath]),
+      { cwd: repoRoot, detached: true, stdio: 'ignore' }
+    );
+  } finally {
+    releaseStartup();
+  }
+  let spawnError: Error | null = null;
+  child.once('error', (error) => { spawnError = error; });
+  child.unref();
+  if (!child.pid) throw new Error('SANDBOX_CONTROL_BROKER_START_FAILED');
+  try {
+    await waitForSandboxControlBrokerStatus({
+      root,
+      statusDir: manifest.publicStatusDir,
+      generation: manifest.generation,
+      pid: child.pid
+    });
+  } catch (error) {
+    child.kill('SIGTERM');
+    if (spawnError) throw spawnError;
+    throw error;
+  }
+}
+
+async function ensureSandboxControlBroker(params: {
   config: SandboxConfig;
   container: string;
   workspace: SandboxWorkspaceIdentity;
-}): void {
+}): Promise<void> {
   const control = sandboxControlPaths({
     base: params.config.controlBase ?? path.join(params.config.home, '.agent-infra', 'sandbox-control'),
     project: params.config.project,
@@ -151,40 +328,61 @@ function ensureSandboxControlBroker(params: {
     identity: params.workspace
   });
   if (!fs.existsSync(control.manifestPath)) return;
-  const manifest = JSON.parse(fs.readFileSync(control.manifestPath, 'utf8')) as { token?: unknown };
+  const validatedManifest = readSandboxControlManifest(control.manifestPath);
+  const containerObservation = await inspectSandboxControlContainer(validatedManifest);
+  if (containerObservation.state === 'unknown') {
+    throw new Error(`SANDBOX_CONTROL_CONTAINER_UNKNOWN: ${containerObservation.reason}`);
+  }
+  if (containerObservation.state === 'absent') {
+    await garbageCollectSandboxControlRoot(control.root, {
+      inspectContainer: (timeoutMs) => inspectSandboxControlContainer(validatedManifest, { timeoutMs })
+    });
+    throw new Error('SANDBOX_CONTROL_CONTAINER_ABSENT: recreate the sandbox container before broker recovery');
+  }
   const brokerPath = path.join(control.root, 'broker.json');
   try {
     const broker = JSON.parse(fs.readFileSync(brokerPath, 'utf8')) as {
       version?: unknown;
       pid?: unknown;
       startTime?: unknown;
+      brokerId?: unknown;
       token?: unknown;
+      generation?: unknown;
     };
     if (
-      broker.version === 1
+      broker.version === 3
       && typeof broker.pid === 'number'
-      && typeof broker.startTime === 'string'
-      && broker.token === manifest.token
-      && processIdentityMatches({ version: 1, pid: broker.pid, startTime: broker.startTime })
-    ) return;
+      && Number.isSafeInteger(broker.pid)
+      && typeof broker.startTime === 'number'
+      && Number.isSafeInteger(broker.startTime)
+      && typeof broker.brokerId === 'string'
+      && broker.brokerId.length > 0
+      && broker.token === validatedManifest.token
+      && broker.generation === validatedManifest.generation
+      && processIdentityMatches({ pid: broker.pid, startTime: broker.startTime })
+    ) {
+      try {
+        const status = readSandboxControlStatus(control.statusDir);
+        if (status.generation === validatedManifest.generation && status.broker.pid === broker.pid
+          && status.broker.startTime === broker.startTime
+          && Date.now() - status.updatedAt <= SANDBOX_CONTROL_STATUS_STALE_MS
+          && status.updatedAt <= Date.now() + SANDBOX_CONTROL_FUTURE_SKEW_MS) return;
+      } catch {
+        // A verified live owner may still be publishing its initial status.
+      }
+      await waitForSandboxControlBrokerStatus({
+        root: control.root,
+        statusDir: control.statusDir,
+        generation: validatedManifest.generation,
+        pid: broker.pid,
+        startTime: broker.startTime
+      });
+      return;
+    }
   } catch {
     // A missing, stale, or malformed broker record is replaced below.
   }
-  fs.rmSync(brokerPath, { force: true });
-  const directory = path.dirname(fileURLToPath(import.meta.url));
-  const extension = path.extname(fileURLToPath(import.meta.url));
-  const internalCli = path.resolve(directory, '..', '..', 'bin', `internal-cli${extension}`);
-  const child = spawn(
-    process.execPath,
-    [internalCli, 'sandbox-control', 'serve', '--manifest', control.manifestPath],
-    { cwd: params.config.repoRoot, detached: true, stdio: 'ignore' }
-  );
-  child.once('error', () => {
-    // Readiness remains fail-closed; a later control request reports the
-    // unavailable broker without turning this detached spawn into an
-    // unhandled process error.
-  });
-  child.unref();
+  await startSandboxControlBroker(params.config.repoRoot, control.manifestPath);
 }
 
 function findingKey(finding: SandboxRecoveryFinding): string {
@@ -194,9 +392,22 @@ function findingKey(finding: SandboxRecoveryFinding): string {
 export function classifySandboxRecovery(snapshot: SandboxRecoverySnapshot): SandboxRecoveryFinding[] {
   const findings: SandboxRecoveryFinding[] = [];
   if (!snapshot.identityOk) {
+    const branchOnly = snapshot.containerIdValid
+      && snapshot.actualBranch === snapshot.expectedBranch
+      && snapshot.expectedWorkspace.mode === 'task-bound'
+      && snapshot.actualWorkspace.mode === 'branch-only';
+    const expectedIdentity = snapshot.expectedWorkspace.mode === 'task-bound'
+      ? `task-bound:${snapshot.expectedWorkspace.taskId}`
+      : snapshot.expectedWorkspace.mode;
+    const actualIdentity = snapshot.actualWorkspace.mode === 'task-bound'
+      ? `task-bound:${snapshot.actualWorkspace.taskId}`
+      : snapshot.actualWorkspace.mode;
     findings.push({
       repairKind: 'hard-failure',
-      message: 'Container identity does not match the requested sandbox branch.'
+      code: branchOnly ? 'SANDBOX_CONTROL_BRANCH_ONLY' : 'SANDBOX_WORKSPACE_IDENTITY_CONFLICT',
+      message: branchOnly
+        ? `Container is branch-only but the requested sandbox is ${expectedIdentity}.`
+        : `Container identity ${actualIdentity} on branch ${snapshot.actualBranch ?? 'unknown'} does not match requested ${expectedIdentity} on branch ${snapshot.expectedBranch}.`
     });
   }
   if (snapshot.runtimeCapabilityOk === false) {
@@ -382,6 +593,7 @@ function expectedMounts(params: {
   const core = sandboxCoreBindMounts(config, branch, {
     workspaceViewRoot: view.root,
     controlDir: control.channelDir,
+    controlStatusDir: control.statusDir,
     ...(params.workspace.mode === 'task-bound'
       ? {
         taskSource: assertSandboxTaskSource(config.repoRoot, params.workspace.taskId),
@@ -603,8 +815,12 @@ export function collectSandboxRecoverySnapshot(params: {
   });
 
   return {
-    identityOk: typeof inspection.Id === 'string'
-      && inspection.Id.length > 0
+    containerIdValid: typeof inspection.Id === 'string' && inspection.Id.length > 0,
+    expectedBranch: params.branch,
+    actualBranch: typeof branchLabel === 'string' ? branchLabel : null,
+    expectedWorkspace: workspace,
+    actualWorkspace: containerWorkspace,
+    identityOk: typeof inspection.Id === 'string' && inspection.Id.length > 0
       && branchLabel === params.branch
       && sameSandboxWorkspaceIdentity(containerWorkspace, workspace),
     runtimeCapabilityOk: enforceCanonicalPlan
@@ -762,7 +978,7 @@ function repairFindings(params: {
 }
 
 function describeFindings(findings: SandboxRecoveryFinding[]): string {
-  return findings.map((finding) => finding.message).join(' ');
+  return findings.map((finding) => finding.code ? `${finding.code}: ${finding.message}` : finding.message).join(' ');
 }
 
 function assess(params: {
@@ -845,13 +1061,19 @@ export async function ensureSandboxReady(params: EnsureSandboxReadyParams): Prom
   const startFn = deps?.start ?? startSandboxContainer;
   const warnings: string[] = [];
   let failure: Error | null = null;
-  ensureSandboxControlBroker({
-    config: params.config,
-    container: params.row.name,
-    workspace: params.workspace ?? { mode: 'branch-only' }
-  });
-
   try {
+    if (deps?.ensureControlBroker) {
+      await deps.ensureControlBroker();
+    } else {
+      await ensureSandboxControlBroker({
+        config: params.config,
+        container: params.row.name,
+        workspace: params.workspace ?? { mode: 'branch-only' }
+      });
+    }
+    if (params.forceRecreate) {
+      throw new Error('Explicit container recreation requested.');
+    }
     if (!params.row.running) {
       startFn(params.engine, params.row.name);
       const initial = assess({
@@ -975,42 +1197,90 @@ export async function ensureSandboxReady(params: EnsureSandboxReadyParams): Prom
 
   const dataBoundary =
     'The existing container writable layer, ordinary /tmp data, processes, tmux sessions, and RAM state may be lost by replacement; worktree and host-managed sandbox data are preserved.';
+  const recoveryTarget = params.workspace?.mode === 'task-bound'
+    ? params.workspace.taskId
+    : params.branch;
   if (!params.allowRecreate || !params.recreate) {
-    throw new Error(`Sandbox recovery failed: ${failure.message} ${dataBoundary} Re-run with --recreate to authorize container-only replacement.`);
+    throw new Error([
+      `Sandbox recovery failed. Run 'ai sandbox start --recreate ${recoveryTarget}' on the host to authorize container-only replacement.`,
+      dataBoundary,
+      `Details: ${failure.message}`
+    ].join('\n'));
+  }
+
+  const existingWorktrees = worktreeDirCandidates(params.config, params.branch).filter((candidate) => fs.existsSync(candidate));
+  if (existingWorktrees.length !== 1) {
+    const remediation = existingWorktrees.length === 0
+      ? `Restore the worktree for '${params.branch}' before recovering the container.`
+      : `Found ${existingWorktrees.map((candidate) => JSON.stringify(candidate)).join(', ')}; `
+        + 'keep the worktree in use, remove the stale directory, then retry.';
+    throw new Error(
+      `SANDBOX_RECOVERY_WORKTREE_SNAPSHOT_INVALID: expected exactly one existing worktree for '${params.branch}', found ${existingWorktrees.length}. ${remediation}`
+    );
+  }
+  const worktreeProbe = worktreeProbeForEngine(params.engine, deps?.probe);
+  const beforeInspection = inspectWorktree(existingWorktrees[0]!, { probe: worktreeProbe });
+  if (beforeInspection.status === 'failed') {
+    throw new Error(`SANDBOX_RECOVERY_WORKTREE_SNAPSHOT_INVALID: ${beforeInspection.message}`);
+  }
+  const before: WorktreeSnapshot = beforeInspection.snapshot;
+  if (before.branch !== params.branch) {
+    throw new Error(
+      `SANDBOX_RECOVERY_WORKTREE_SNAPSHOT_INVALID: expected branch '${params.branch}', found '${before.branch}'. `
+      + `Run 'git -C ${JSON.stringify(before.worktree)} checkout ${params.branch}' to restore the expected branch, then retry.`
+    );
   }
 
-  const warning = `Sandbox recovery failed in place. Replacing only the container. ${dataBoundary}`;
+  const warning = params.forceRecreate
+    ? `Explicit sandbox recreation requested. Replacing only the container. ${dataBoundary}`
+    : `Sandbox recovery failed in place. Replacing only the container. ${dataBoundary}`;
   warnings.push(warning);
   (params.writeWarning ?? ((message) => process.stderr.write(`${message}\n`)))(warning);
   const runVerboseFn = deps?.runVerbose ?? runVerboseEngine;
-  if (params.row.running) {
-    runVerboseFn(params.engine, 'docker', ['stop', params.row.name]);
-  }
-  runVerboseFn(params.engine, 'docker', ['rm', params.row.name]);
-  await params.recreate(params.branch);
+  let replacementResult: SandboxReadyResult | null = null;
+  let replacementFailure: unknown = null;
+  try {
+    if (params.row.running) {
+      runVerboseFn(params.engine, 'docker', ['stop', params.row.name]);
+    }
+    runVerboseFn(params.engine, 'docker', ['rm', params.row.name]);
+    await params.recreate(params.branch);
 
-  const rows = (deps?.fetchRows ?? fetchSandboxRows)(
-    params.engine,
-    sandboxLabel(params.config),
-    sandboxBranchLabel(params.config)
-  );
-  const replacement = selectSandboxContainer(
-    [...rows.running, ...rows.nonRunning],
-    containerNameCandidates(params.config, params.branch)
-  );
-  if (!replacement?.running) {
-    throw new Error('Replacement sandbox container was not found in a running state.');
+    const rows = (deps?.fetchRows ?? fetchSandboxRows)(
+      params.engine,
+      sandboxLabel(params.config),
+      sandboxBranchLabel(params.config)
+    );
+    const replacement = selectSandboxContainer(
+      [...rows.running, ...rows.nonRunning],
+      containerNameCandidates(params.config, params.branch)
+    );
+    if (!replacement?.running) {
+      throw new Error('Replacement sandbox container was not found in a running state.');
+    }
+    const final = assess({
+      config: params.config,
+      engine: params.engine,
+      branch: params.branch,
+      workspace: params.workspace,
+      container: replacement.name,
+      deps
+    });
+    if (final.findings.length > 0) {
+      throw new Error(`Replacement sandbox readiness check failed: ${describeFindings(final.findings)}`);
+    }
+    replacementResult = { container: replacement.name, path: 'recreated', warnings };
+  } catch (error) {
+    replacementFailure = error;
   }
-  const final = assess({
-    config: params.config,
-    engine: params.engine,
-    branch: params.branch,
-    workspace: params.workspace,
-    container: replacement.name,
-    deps
-  });
-  if (final.findings.length > 0) {
-    throw new Error(`Replacement sandbox readiness check failed: ${describeFindings(final.findings)}`);
+
+  const afterInspection = inspectWorktree(before.worktree, { probe: worktreeProbe });
+  if (afterInspection.status === 'failed' || afterInspection.snapshot.identity !== before.identity) {
+    const detail = afterInspection.status === 'failed'
+      ? afterInspection.message
+      : `snapshot changed from ${before.identity} to ${afterInspection.snapshot.identity}`;
+    throw new Error(`SANDBOX_RECOVERY_WORKTREE_CHANGED: ${detail}`, replacementFailure === null ? undefined : { cause: replacementFailure });
   }
-  return { container: replacement.name, path: 'recreated', warnings };
+  if (replacementFailure !== null) throw replacementFailure;
+  return replacementResult!;
 }