@fitlab-ai/agent-infra 0.9.4 → 0.9.6

package/lib/sandbox/mounts.ts

@@ -6,6 +6,7 @@ import {
   worktreeDirCandidates
 } from './constants.ts';
 import type { SandboxConfig } from './config.ts';
+import { sandboxWorkspaceViewStatePaths } from './workspace-view.ts';
 
 export type SandboxBindMountDeclaration = {
   hostPaths: string[];
@@ -24,10 +25,25 @@ export function sandboxCoreBindMounts(
     shellConfigHostDir?: string;
     workspaceViewRoot: string;
     controlDir: string;
+    controlStatusDir: string;
     taskSource?: string;
     taskId?: string;
   }
 ): SandboxBindMountDeclaration[] {
+  const taskBound = Boolean(overrides.taskSource && overrides.taskId);
+  const workspaceMounts = sandboxWorkspaceViewStatePaths(overrides.workspaceViewRoot).map(
+    ({ state, hostPath }) => taskBound && state === 'active'
+      ? {
+          hostPaths: [path.join(hostPath, '.short-ids.json')],
+          containerPath: '/workspace/.agents/workspace/active/.short-ids.json',
+          readOnly: true
+        }
+      : {
+          hostPaths: [hostPath],
+          containerPath: path.posix.join('/workspace/.agents/workspace', state),
+          readOnly: true
+        }
+  );
   const mounts: SandboxBindMountDeclaration[] = [
     {
       hostPaths: overrides.worktree
@@ -36,11 +52,16 @@ export function sandboxCoreBindMounts(
       containerPath: '/workspace',
       readOnly: false
     },
-    {
-      hostPaths: [overrides.workspaceViewRoot],
-      containerPath: '/workspace/.agents/workspace',
-      readOnly: true
-    },
+    ...workspaceMounts,
+  ];
+  if (taskBound) {
+    mounts.push({
+      hostPaths: [overrides.taskSource!],
+      containerPath: path.posix.join('/workspace/.agents/workspace/active', overrides.taskId!),
+      readOnly: false
+    });
+  }
+  mounts.push(
     {
       hostPaths: [shareCommonDir(config)],
       containerPath: '/share/common',
@@ -58,15 +79,12 @@ export function sandboxCoreBindMounts(
       containerPath: '/home/devuser/.host-shell-config',
       readOnly: true
     }
-  ];
-  if (overrides.taskSource && overrides.taskId) {
-    mounts.splice(2, 0, {
-      hostPaths: [overrides.taskSource],
-      containerPath: path.posix.join('/workspace/.agents/workspace/active', overrides.taskId),
-      readOnly: false
-    });
-  }
+  );
   mounts.push({
+    hostPaths: [overrides.controlStatusDir],
+    containerPath: '/run/agent-infra/control-status',
+    readOnly: true
+  }, {
     hostPaths: [overrides.controlDir],
     containerPath: '/run/agent-infra/control',
     readOnly: false

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,8 +31,10 @@ import {
 import {
   parseSandboxWorkspaceIdentity,
   sameSandboxWorkspaceIdentity,
+  type SandboxContainerWorkspaceIdentity,
   type SandboxWorkspaceIdentity
 } from './workspace-identity.ts';
+import { inspectWorktree, type WorktreeSnapshot } from './worktree-safety.ts';
 import { runEngine, runOkEngine, runVerboseEngine } from './shell.ts';
 import {
   declaredTmpfsSeedEntries,
@@ -45,10 +49,24 @@ 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,
+  isSandboxControlRootQuiescing
+} from './control/lifecycle.ts';
+import {
+  appendSandboxControlAudit,
+  readSandboxControlStatus
+} from './control/state.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 +84,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,6 +141,7 @@ type RecoveryCommandDeps = {
   runVerbose?: typeof runVerboseEngine;
   start?: typeof startSandboxContainer;
   fetchRows?: typeof fetchSandboxRows;
+  ensureControlBroker?: () => void | Promise<void>;
 };
 
 type EnsureSandboxReadyParams = {
@@ -139,11 +163,153 @@ type ExpectedMount = {
   expectedRW: boolean;
 };
 
-function ensureSandboxControlBroker(params: {
+async function waitForSandboxControlBrokerStatus(params: {
+  root: string;
+  statusDir: string;
+  generation: string;
+  pid: number;
+  startTime?: string;
+}, 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 = JSON.parse(fs.readFileSync(manifestPath, 'utf8')) as SandboxControlManifest;
+  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;
+      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 && typeof broker.pid === 'number' && typeof broker.startTime === 'string'
+      ? { version: 1 as const, pid: broker.pid, startTime: broker.startTime }
+      : null;
+    const live = owner !== null && processIdentityMatches(owner);
+    if (
+      live
+      && owner !== null
+      && broker !== null
+      && broker.version === 2
+      && broker.token === manifest.token
+      && broker.generation === manifest.generation
+    ) {
+      await waitForSandboxControlBrokerStatus({
+        root,
+        statusDir: manifest.publicStatusDir,
+        generation: manifest.generation,
+        pid: owner.pid,
+        startTime: owner.startTime
+      });
+      return;
+    }
+    if (live && owner) {
+      const deadline = Date.now() + 5_000;
+      while (Date.now() < deadline) {
+        await new Promise((resolve) => setTimeout(resolve, 25));
+        try {
+          if (fs.readFileSync(brokerPath, 'utf8') !== brokerSnapshot) {
+            return startSandboxControlBroker(repoRoot, manifestPath);
+          }
+        } catch (error) {
+          if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
+            return startSandboxControlBroker(repoRoot, manifestPath);
+          }
+          throw error;
+        }
+        if (!processIdentityMatches(owner)) {
+          removePidFileIfMatches(brokerPath, brokerSnapshot);
+          return startSandboxControlBroker(repoRoot, manifestPath);
+        }
+      }
+      throw new Error('SANDBOX_CONTROL_BROKER_OWNER_ACTIVE');
+    }
+  }
+  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,7 +317,10 @@ 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 manifest = JSON.parse(fs.readFileSync(control.manifestPath, 'utf8')) as SandboxControlManifest;
+  if (manifest.version !== 3) {
+    throw new Error('SANDBOX_CONTROL_MANIFEST_VERSION_INVALID: expected version 3; container-only recreation is required');
+  }
   const brokerPath = path.join(control.root, 'broker.json');
   try {
     const broker = JSON.parse(fs.readFileSync(brokerPath, 'utf8')) as {
@@ -159,32 +328,38 @@ function ensureSandboxControlBroker(params: {
       pid?: unknown;
       startTime?: unknown;
       token?: unknown;
+      generation?: unknown;
     };
     if (
-      broker.version === 1
+      broker.version === 2
       && typeof broker.pid === 'number'
       && typeof broker.startTime === 'string'
       && broker.token === manifest.token
+      && broker.generation === manifest.generation
       && processIdentityMatches({ version: 1, pid: broker.pid, startTime: broker.startTime })
-    ) return;
+    ) {
+      try {
+        const status = readSandboxControlStatus(control.statusDir);
+        if (status.generation === manifest.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: manifest.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 +369,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 +570,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 +792,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 +955,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 +1038,16 @@ 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.row.running) {
       startFn(params.engine, params.row.name);
       const initial = assess({
@@ -975,42 +1171,87 @@ 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 beforeInspection = inspectWorktree(existingWorktrees[0]!);
+  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}`;
   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);
+  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!;
 }