""" workspace.py - Fixture-seeded workspace and fresh caller-session materialization. Provides deterministic, standard-library-only workspace preparation beneath an already allocated empty attempt root. """ from __future__ import annotations import datetime import hashlib import json import os import posixpath import re import secrets import shutil import subprocess from dataclasses import dataclass from pathlib import Path from typing import Any, Optional from scripts.agent_benchmark.manifest import ( AssetMapping, Manifest, digest_workspace_inputs, ) RUN_ID_RE = re.compile(r"^run-\d{8}T\d{6}Z-[0-9a-f]{12}$") CELL_ID_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$") class WorkspaceError(Exception): """Base exception for workspace preparation failures.""" class WorkspaceValidationError(WorkspaceError): """Raised when identity or manifest configuration is invalid for workspace preparation.""" class WorkspacePathError(WorkspaceError): """Raised when an attempt root or asset path violates containment, symlink, or collision rules.""" class WorkspaceChecksumError(WorkspaceError): """Raised when initial fixture or materialized workspace checksum fails verification.""" class TestbedError(WorkspaceError): """Raised when the testbed provenance check fails or testbed is dirty/mutated.""" @dataclass(frozen=True) class AttemptIdentity: run_id: str cell_id: str repetition: int attempt: int class _PreparationOwnership: """Tracks filesystem paths created and owned by a workspace preparation transaction.""" def __init__(self, attempt_root: Path) -> None: self.attempt_root = attempt_root.resolve() self.created_paths: list[Path] = [] def register(self, path: Path) -> None: resolved = path.resolve() if resolved not in self.created_paths: self.created_paths.append(resolved) def unregister(self, path: Path) -> None: resolved = path.resolve() if resolved in self.created_paths: self.created_paths.remove(resolved) def rollback(self) -> None: """Remove all registered paths created by this invocation in reverse order.""" for path in reversed(self.created_paths): if not path.exists() and not path.is_symlink(): continue try: if path.is_dir() and not path.is_symlink(): shutil.rmtree(path, ignore_errors=True) else: path.unlink(missing_ok=True) except OSError: pass def _destinations_conflict(path1: str, path2: str) -> bool: """Return True if path1 and path2 conflict as identical or ancestor/descendant paths.""" if path1 == path2: return True return path2.startswith(path1 + "/") or path1.startswith(path2 + "/") @dataclass(frozen=True) class TestbedProvenance: path: str branch: str head: str status_digest: str clean: bool @dataclass(frozen=True) class PreparedWorkspace: identity: AttemptIdentity attempt_root: str workspace_dir: str session_dir: str session_id: str session_is_fresh: bool workspace_checksum: str setup_cache_policy: str testbed_provenance: TestbedProvenance prepared_at: str def _find_repo_root(start_path: Path) -> Path: """Find repository root by searching upward for Makefile or .git.""" candidate = start_path.resolve() for _ in range(20): if (candidate / "Makefile").is_file() or (candidate / ".git").exists(): return candidate parent = candidate.parent if parent == candidate: break candidate = parent return Path.cwd().resolve() def validate_attempt_identity( identity: AttemptIdentity, manifest: Manifest | None = None ) -> None: """Validate format and bounds of an AttemptIdentity. Raises: WorkspaceValidationError: If any field fails format or boundary checks. """ if not isinstance(identity, AttemptIdentity): raise WorkspaceValidationError("identity must be an AttemptIdentity instance") if not isinstance(identity.run_id, str) or not RUN_ID_RE.match(identity.run_id): raise WorkspaceValidationError(f"invalid run_id format '{identity.run_id}'") if not isinstance(identity.cell_id, str) or not CELL_ID_RE.match(identity.cell_id): raise WorkspaceValidationError(f"invalid cell_id format '{identity.cell_id}'") if ( isinstance(identity.repetition, bool) or not isinstance(identity.repetition, int) or identity.repetition < 1 ): raise WorkspaceValidationError("repetition must be a positive integer >= 1") if ( isinstance(identity.attempt, bool) or not isinstance(identity.attempt, int) or identity.attempt < 1 ): raise WorkspaceValidationError("attempt must be a positive integer >= 1") if manifest is not None: valid_cell_ids = {cell.id for cell in manifest.matrix} if identity.cell_id not in valid_cell_ids: raise WorkspaceValidationError( f"cell_id '{identity.cell_id}' not found in manifest matrix" ) if identity.repetition > manifest.repetitions: raise WorkspaceValidationError( f"repetition {identity.repetition} exceeds manifest repetitions ({manifest.repetitions})" ) def inspect_testbed_provenance(testbed_path: str | Path) -> TestbedProvenance: """Capture Git provenance of the runtime testbed directory. Requires a clean Git working copy. Never modifies the repository. Raises: TestbedError: If directory does not exist, is not a git repo, or is dirty. """ path = Path(testbed_path).resolve() if not path.exists() or not path.is_dir(): raise TestbedError( f"testbed directory '{path}' does not exist or is not a directory" ) try: proc_status = subprocess.run( ["git", "status", "--porcelain=v1", "--untracked-files=all"], cwd=path, capture_output=True, text=True, check=False, ) except Exception as exc: raise TestbedError(f"failed to run git status in testbed '{path}': {exc}") if proc_status.returncode != 0: raise TestbedError(f"git status returned non-zero exit code in testbed '{path}'") status_output = proc_status.stdout.strip() if status_output: raise TestbedError(f"testbed repository '{path}' is dirty:\n{status_output}") try: proc_branch = subprocess.run( ["git", "branch", "--show-current"], cwd=path, capture_output=True, text=True, check=False, ) branch = proc_branch.stdout.strip() if not branch: proc_head_ref = subprocess.run( ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=path, capture_output=True, text=True, check=False, ) branch = proc_head_ref.stdout.strip() except Exception as exc: raise TestbedError(f"failed to inspect git branch in testbed '{path}': {exc}") try: proc_head = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=path, capture_output=True, text=True, check=False, ) except Exception as exc: raise TestbedError(f"failed to inspect git HEAD in testbed '{path}': {exc}") if proc_head.returncode != 0: raise TestbedError(f"git rev-parse HEAD failed in testbed '{path}'") head = proc_head.stdout.strip() status_raw = f"branch:{branch}\nhead:{head}\nstatus:{status_output}".encode("utf-8") status_digest = "sha256:" + hashlib.sha256(status_raw).hexdigest() return TestbedProvenance( path=str(path), branch=branch, head=head, status_digest=status_digest, clean=True, ) def prepare_workspace( manifest: Manifest, attempt_root: str | Path, identity: AttemptIdentity, repo_root: str | Path | None = None, ) -> PreparedWorkspace: """Materialize clean workspace and session under an allocated empty attempt root. Args: manifest: Validated benchmark manifest object. attempt_root: Path to an exclusively allocated empty attempt root directory. identity: Frozen AttemptIdentity specifying run/cell/repetition/attempt. repo_root: Optional repository root path for resolving testbed/assets. Returns: PreparedWorkspace dataclass containing materialized metadata. Raises: WorkspaceValidationError: If identity or manifest schema/policy checks fail. WorkspacePathError: If path containment, symlink, or root canonicality fails. WorkspaceChecksumError: If fixture or workspace content digest verification fails. TestbedError: If testbed provenance inspection fails or testbed is dirty/modified. """ attempt_root_input = Path(attempt_root) if repo_root is None: repo_root_path = _find_repo_root(attempt_root_input) else: repo_root_path = Path(repo_root).resolve() # Phase 1: Validate complete plan before any visible publication. validated_inputs = _validate_complete_plan( manifest, attempt_root_input, identity, repo_root_path ) resolved_attempt_root = validated_inputs["resolved_attempt_root"] owned = _PreparationOwnership(resolved_attempt_root) # Phase 2: Stage preparation beneath the validated empty attempt root in a private owned directory. # Phase 3: Verify staged content (checksum, postflight). # Phase 4: Publish metadata and prepared.json. # All three phases are bounded; any exception triggers rollback of function-owned artifacts. try: staging = _stage_preparation(validated_inputs, owned) _verify_staged_preparation(staging, manifest) prepared = _publish_preparation(staging, owned) except Exception: _rollback_owned_preparation(owned) raise return prepared def _validate_complete_plan( manifest: Manifest, attempt_root_input: Path, identity: AttemptIdentity, repo_root_path: Path, ) -> dict[str, Any]: """Validate the entire preparation plan without creating any files. Returns a dict of resolved paths and validated state for staging. """ # 1. Validate identity against manifest validate_attempt_identity(identity, manifest) if manifest.session_policy != "fresh": raise WorkspaceValidationError( f"unsupported session_policy '{manifest.session_policy}', expected 'fresh'" ) if manifest.setup_cache_policy != "isolated": raise WorkspaceValidationError( f"unsupported setup_cache_policy '{manifest.setup_cache_policy}', expected 'isolated'" ) # 2. Path & Symlink validation on attempt_root raw_path = attempt_root_input if raw_path.is_symlink(): raise WorkspacePathError("attempt_root cannot be a symlink") # Walk up parent segments to ensure no symlinks in path curr = raw_path while curr != curr.parent: if curr.is_symlink(): raise WorkspacePathError(f"path segment '{curr}' in attempt_root is a symlink") if curr == repo_root_path: break curr = curr.parent resolved_attempt_root = raw_path.resolve() # Verify expected canonical relative path rep_segment = f"repetition-{identity.repetition:04d}" att_segment = f"attempt-{identity.attempt:06d}" expected_rel = ( Path(manifest.output_root) / identity.run_id / "cells" / identity.cell_id / rep_segment / att_segment ) expected_abs = (repo_root_path / expected_rel).resolve() if resolved_attempt_root != expected_abs: raise WorkspacePathError( f"attempt_root '{resolved_attempt_root}' does not match expected canonical path '{expected_abs}'" ) if not resolved_attempt_root.exists() or not resolved_attempt_root.is_dir(): raise WorkspacePathError( f"attempt_root '{resolved_attempt_root}' does not exist or is not a directory" ) # Verify attempt_root is empty if list(resolved_attempt_root.iterdir()): raise WorkspacePathError(f"attempt_root '{resolved_attempt_root}' is not empty") # 3. Testbed preflight provenance check (must not be beneath attempt_root) testbed_path = (repo_root_path / manifest.testbed).resolve() try: testbed_path.relative_to(resolved_attempt_root) raise WorkspacePathError("testbed cannot be located beneath attempt_root") except ValueError: pass testbed_before = inspect_testbed_provenance(testbed_path) # 4. Check fixture checksum before copying computed_fixture_checksum = digest_workspace_inputs(manifest.fixture.assets) if computed_fixture_checksum != manifest.fixture.checksum: raise WorkspaceChecksumError( f"declared fixture checksum '{manifest.fixture.checksum}' does not match computed fixture asset digest '{computed_fixture_checksum}'" ) # 5. Validate all asset sources exist, aren't symlinks, are within repo root # and check for destination collisions before any creation. workspace_dir = resolved_attempt_root / "workspace" if workspace_dir.exists(): raise WorkspacePathError("workspace child directory already exists") asset_validations: list[dict[str, Any]] = [] workspace_destinations: set[str] = set() for asset in manifest.fixture.assets: raw_src_path = repo_root_path / asset.source if raw_src_path.is_symlink(): raise WorkspacePathError(f"asset source '{asset.source}' cannot be a symlink") src_path = raw_src_path.resolve() if src_path.is_symlink(): raise WorkspacePathError(f"asset source '{asset.source}' cannot be a symlink") if not src_path.exists() or not src_path.is_file(): raise WorkspacePathError( f"asset source '{asset.source}' does not exist or is not a regular file" ) try: src_path.relative_to(repo_root_path) except ValueError: raise WorkspacePathError(f"asset source '{asset.source}' escapes repository root") wp = asset.workspace_path if "\\" in wp or ":" in wp: raise WorkspacePathError(f"asset workspace_path '{wp}' contains invalid characters") wp_norm = posixpath.normpath(wp) if wp != wp_norm or wp.startswith("/") or wp.startswith(".."): raise WorkspacePathError( f"asset workspace_path '{wp}' is not in canonical posix relative form" ) target_file = workspace_dir / wp_norm try: target_file.relative_to(workspace_dir) except ValueError: raise WorkspacePathError(f"asset workspace_path '{wp}' escapes workspace directory") # Check for destination collisions before copying if target_file.exists(): raise WorkspacePathError( f"asset destination '{wp}' already exists in workspace" ) # Check parent path for symlinks parent_check = target_file.parent while parent_check != workspace_dir: if parent_check.is_symlink(): raise WorkspacePathError( f"symlink detected in asset target path '{parent_check}'" ) parent_check = parent_check.parent for existing in workspace_destinations: if _destinations_conflict(existing, wp_norm): raise WorkspacePathError( f"asset workspace_path '{wp_norm}' conflicts with '{existing}'" ) workspace_destinations.add(wp_norm) asset_validations.append({ "source": asset.source, "src_path": src_path, "workspace_path": wp_norm, "target_path": target_file, "content": asset.content, }) # 6. Verify prompt is not copied unless declared as an asset prompt_declared = any( a.source == manifest.fixture.prompt or a.workspace_path == manifest.fixture.prompt for a in manifest.fixture.assets ) if not prompt_declared: prompt_in_ws = workspace_dir / manifest.fixture.prompt if prompt_in_ws.exists(): raise WorkspacePathError( "prompt file materialized in workspace without being declared as an asset" ) return { "resolved_attempt_root": resolved_attempt_root, "workspace_dir": workspace_dir, "testbed_path": testbed_path, "testbed_before": testbed_before, "asset_validations": asset_validations, "prompt_declared": prompt_declared, "identity": identity, } def _stage_preparation( validated_inputs: dict[str, Any], owned: _PreparationOwnership ) -> dict[str, Any]: """Create workspace/ and session/ directories and copy assets in private staging. This is the staging phase. Any failure here triggers rollback of owned entries. """ resolved_attempt_root = validated_inputs["resolved_attempt_root"] staging_dir = resolved_attempt_root / f".staging-{secrets.token_hex(6)}" staging_dir.mkdir(parents=False, exist_ok=False) owned.register(staging_dir) staging_workspace_dir = staging_dir / "workspace" staging_workspace_dir.mkdir(parents=False, exist_ok=False) materialized_assets: list[AssetMapping] = [] for av in validated_inputs["asset_validations"]: src_path = av["src_path"] target_file = staging_workspace_dir / av["workspace_path"] content = src_path.read_bytes() if av["content"] and av["content"] != content: raise WorkspaceChecksumError( f"asset content for '{av['source']}' does not match resolved file content" ) target_file.parent.mkdir(parents=True, exist_ok=True) target_file.write_bytes(content) materialized_assets.append( AssetMapping(source=av["source"], workspace_path=av["workspace_path"], content=content) ) staging_session_dir = staging_dir / "session" staging_session_dir.mkdir(parents=False, exist_ok=False) return { "resolved_attempt_root": resolved_attempt_root, "staging_dir": staging_dir, "workspace_dir": staging_workspace_dir, "session_dir": staging_session_dir, "materialized_assets": materialized_assets, "validated_inputs": validated_inputs, } def _verify_staged_preparation( staging: dict[str, Any], manifest: Manifest ) -> None: """Verify staged workspace checksum and testbed postflight.""" workspace_dir = staging["workspace_dir"] testbed_path = staging["validated_inputs"]["testbed_path"] testbed_before = staging["validated_inputs"]["testbed_before"] # Recompute workspace checksum after materialization actual_assets: list[AssetMapping] = [] for root, _, files in os.walk(workspace_dir): for f in files: fp = Path(root) / f if fp.is_symlink(): raise WorkspacePathError(f"symlink created in workspace directory '{fp}'") rel_wp = fp.relative_to(workspace_dir).as_posix() actual_assets.append( AssetMapping(source="", workspace_path=rel_wp, content=fp.read_bytes()) ) computed_workspace_checksum = digest_workspace_inputs(actual_assets) if computed_workspace_checksum != manifest.fixture.checksum: raise WorkspaceChecksumError( f"materialized workspace checksum '{computed_workspace_checksum}' does not match declared fixture checksum '{manifest.fixture.checksum}'" ) # Postflight testbed check testbed_after = inspect_testbed_provenance(testbed_path) if testbed_before != testbed_after: raise TestbedError("testbed repository state was modified during workspace preparation") # Store testbed_after and workspace_checksum in staging for publish phase staging["testbed_after"] = testbed_after staging["workspace_checksum"] = computed_workspace_checksum def _publish_owned_directory( staging_dir: Path, final_dir: Path, owned: _PreparationOwnership, collision_message: str, ) -> None: """Exclusively create final_dir, register ownership, and move staged contents into it.""" try: final_dir.mkdir(exist_ok=False) except (FileExistsError, OSError): raise WorkspacePathError(collision_message) owned.register(final_dir) for item in staging_dir.iterdir(): item.rename(final_dir / item.name) def _publish_preparation( staging: dict[str, Any], owned: _PreparationOwnership ) -> PreparedWorkspace: """Generate session ID, build metadata, publish final entries, and write prepared.json.""" validated_inputs = staging["validated_inputs"] identity = validated_inputs["identity"] token = secrets.token_hex(6) session_id = ( f"session-{identity.run_id}-{identity.cell_id}-" f"rep{identity.repetition:04d}-att{identity.attempt:06d}-{token}" ) session_is_fresh = True staging_dir = staging["staging_dir"] staging_workspace_dir = staging["workspace_dir"] staging_session_dir = staging["session_dir"] resolved_attempt_root = staging["resolved_attempt_root"] testbed_after = staging["testbed_after"] workspace_checksum = staging["workspace_checksum"] final_workspace_dir = resolved_attempt_root / "workspace" _publish_owned_directory( staging_workspace_dir, final_workspace_dir, owned, "workspace child directory already exists in attempt_root", ) final_session_dir = resolved_attempt_root / "session" _publish_owned_directory( staging_session_dir, final_session_dir, owned, "session child directory already exists in attempt_root", ) if staging_dir.exists(): shutil.rmtree(staging_dir, ignore_errors=True) owned.unregister(staging_dir) prepared_at = datetime.datetime.now(datetime.timezone.utc).isoformat() prepared = PreparedWorkspace( identity=identity, attempt_root=str(resolved_attempt_root), workspace_dir=str(final_workspace_dir), session_dir=str(final_session_dir), session_id=session_id, session_is_fresh=session_is_fresh, workspace_checksum=workspace_checksum, setup_cache_policy="isolated", testbed_provenance=testbed_after, prepared_at=prepared_at, ) prepared_data = { "identity": { "run_id": identity.run_id, "cell_id": identity.cell_id, "repetition": identity.repetition, "attempt": identity.attempt, }, "attempt_root": prepared.attempt_root, "workspace_dir": prepared.workspace_dir, "session_dir": prepared.session_dir, "session_id": prepared.session_id, "session_is_fresh": prepared.session_is_fresh, "workspace_checksum": prepared.workspace_checksum, "setup_cache_policy": prepared.setup_cache_policy, "testbed_provenance": { "path": prepared.testbed_provenance.path, "branch": prepared.testbed_provenance.branch, "head": prepared.testbed_provenance.head, "status_digest": prepared.testbed_provenance.status_digest, "clean": prepared.testbed_provenance.clean, }, "prepared_at": prepared.prepared_at, } prepared_json_path = resolved_attempt_root / "prepared.json" if prepared_json_path.exists() or prepared_json_path.is_symlink(): raise WorkspacePathError("prepared.json already exists in attempt_root") try: with prepared_json_path.open("x", encoding="utf-8") as f: json.dump(prepared_data, f, indent=2) except FileExistsError: raise WorkspacePathError("prepared.json already exists in attempt_root") owned.register(prepared_json_path) return prepared def _rollback_owned_preparation(owned: _PreparationOwnership) -> None: """Remove paths explicitly registered as created by this preparation transaction.""" owned.rollback()