"""Private, atomic checkpoints for host-executed GEWU research runs.""" from __future__ import annotations import argparse import importlib import json import math import os import re import secrets import stat import sys from datetime import UTC, datetime from pathlib import Path from types import ModuleType from typing import Any if __package__: manifest_module: ModuleType = importlib.import_module(".manifest_tool", __package__) else: sys.path.insert(0, str(Path(__file__).resolve().parent)) manifest_module = importlib.import_module("manifest_tool") load_manifest = manifest_module.load_manifest verify = manifest_module.verify _RUN_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{1,63}$") _SHA256 = re.compile(r"^[a-f0-9]{64}$") _MUTABLE = { "stage", "sources", "coverage", "budgets", "pending_operations", "proposals", "outputs", "last_failure", "runtime", } _FORBIDDEN_KEYS = { "token", "secret", "password", "credential", "source_body", "content", "cookie", "authorization", "auth_header", } def _reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: result: dict[str, Any] = {} for key, value in pairs: if key in result: raise ValueError("checkpoint JSON contains a duplicate key") result[key] = value return result def _safe_payload(value: Any) -> None: if isinstance(value, dict): for key, item in value.items(): normalized = str(key).lower().replace("-", "_") if ( normalized in _FORBIDDEN_KEYS or any( part in normalized for part in ("secret", "password", "credential", "bearer", "api_key") ) or normalized.endswith("_token") ): raise ValueError("checkpoint cannot store credentials or source bodies") _safe_payload(item) elif isinstance(value, list): for item in value: _safe_payload(item) elif isinstance(value, float) and not math.isfinite(value): raise ValueError("checkpoint contains a non-finite number") elif not isinstance(value, (str, int, float, bool, type(None))): raise ValueError("checkpoint contains a non-JSON value") def _runs_directory(workspace: Path) -> Path: root = workspace.expanduser().absolute() if root.is_symlink(): raise ValueError("workspace root must not be a symbolic link") runs = root / "Runs" if runs.is_symlink(): raise ValueError("Runs directory must not be a symbolic link") runs.mkdir(parents=True, exist_ok=True, mode=0o700) if not runs.is_dir() or runs.is_symlink(): raise ValueError("Runs path is unavailable") if os.name != "nt": runs.chmod(stat.S_IRWXU) return runs def _checkpoint_path(workspace: Path, run_id: str) -> Path: if not _RUN_ID.fullmatch(run_id): raise ValueError("run ID is invalid") return _runs_directory(workspace) / f"{run_id}.json" def _write_atomic(path: Path, value: dict[str, Any]) -> None: _safe_payload(value) payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) if len(payload.encode("utf-8")) > 1_000_000: raise ValueError("checkpoint exceeds the one-megabyte limit") temporary = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp") try: descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as stream: stream.write(payload + "\n") stream.flush() os.fsync(stream.fileno()) os.replace(temporary, path) if os.name != "nt": path.chmod(stat.S_IRUSR | stat.S_IWUSR) directory_fd = os.open(path.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) finally: temporary.unlink(missing_ok=True) def _module_pins(manifest: dict[str, Any]) -> list[dict[str, str]]: return [ { "module_id": module["module_id"], "entry": module["entry"], "contract_id": module["provided_contract"]["id"], "contract_version": module["provided_contract"]["version"], "implementation_version": module["implementation_version"], "sha256": module["sha256"], } for module in manifest["modules"] ] def create_checkpoint( workspace: Path, *, run_id: str, objective: str, projects: list[str], include_shared: bool, lens: str, budgets: dict[str, int] | None = None, ) -> dict[str, Any]: manifest = load_manifest() verify(manifest) if not objective.strip() or len(objective) > 4_000: raise ValueError("objective must contain 1 to 4000 characters") if ( not projects or len(projects) > 100 or len(set(projects)) != len(projects) or any( not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", project) for project in projects ) ): raise ValueError("scope must contain distinct explicit project IDs") if not lens.strip() or len(lens) > 2_000: raise ValueError("lens must contain 1 to 2000 characters") path = _checkpoint_path(workspace, run_id) if path.is_symlink() or path.exists(): raise ValueError("run checkpoint already exists") now = datetime.now(UTC).isoformat() checkpoint: dict[str, Any] = { "checkpoint_version": "0.1.0", "run_id": run_id, "created_at": now, "updated_at": now, "pack": { "pack_id": manifest["pack_id"], "pack_version": manifest["pack_version"], "snapshot_digest": manifest["snapshot"]["digest"], "modules": _module_pins(manifest), "service_contracts": [ {"contract_id": item["contract_id"], "sha256": item["sha256"]} for item in manifest["verified_service_contracts"] ], }, "objective": objective, "scope": {"projects": projects, "include_shared": include_shared}, "lens": lens, "stage": "context", "sources": [], "coverage": {"selected": [], "included": [], "excluded": [], "skipped": []}, "budgets": budgets or {}, "runtime": {"host": None, "model": None}, "pending_operations": [], "proposals": [], "outputs": [], "last_failure": None, } _validate_checkpoint(checkpoint) _write_atomic(path, checkpoint) return checkpoint def load_checkpoint(workspace: Path, run_id: str) -> dict[str, Any]: manifest = load_manifest() verify(manifest) path = _checkpoint_path(workspace, run_id) if path.is_symlink() or not path.is_file() or path.stat().st_size > 1_000_000: raise ValueError("run checkpoint is unavailable") checkpoint = json.loads( path.read_text(encoding="utf-8"), object_pairs_hook=_reject_duplicate_keys ) if not isinstance(checkpoint, dict) or checkpoint.get("run_id") != run_id: raise ValueError("run checkpoint identity is invalid") pinned_pack = checkpoint.get("pack") if not isinstance(pinned_pack, dict) or ( pinned_pack.get("snapshot_digest") != manifest["snapshot"]["digest"] or pinned_pack.get("modules") != _module_pins(manifest) ): raise ValueError("pack changed since this run; revalidate before resuming") _safe_payload(checkpoint) _validate_checkpoint(checkpoint) return checkpoint def update_checkpoint(workspace: Path, run_id: str, patch: dict[str, Any]) -> dict[str, Any]: checkpoint = load_checkpoint(workspace, run_id) if not patch or not set(patch).issubset(_MUTABLE): raise ValueError("checkpoint update contains unsupported fields") _safe_payload(patch) updated = {**checkpoint, **patch, "updated_at": datetime.now(UTC).isoformat()} _validate_checkpoint(updated) _write_atomic(_checkpoint_path(workspace, run_id), updated) return updated def _validate_checkpoint(checkpoint: dict[str, Any]) -> None: required = { "checkpoint_version", "run_id", "created_at", "updated_at", "pack", "objective", "scope", "lens", "stage", "sources", "coverage", "budgets", "runtime", "pending_operations", "proposals", "outputs", "last_failure", } if set(checkpoint) != required or checkpoint.get("checkpoint_version") != "0.1.0": raise ValueError("checkpoint fields or version are invalid") if not isinstance(checkpoint["objective"], str) or not checkpoint["objective"].strip(): raise ValueError("checkpoint objective is required") if not isinstance(checkpoint["lens"], str) or not checkpoint["lens"].strip(): raise ValueError("checkpoint lens is required") scope = checkpoint["scope"] if ( not isinstance(scope, dict) or set(scope) != {"projects", "include_shared"} or not isinstance(scope["projects"], list) or not scope["projects"] or len(scope["projects"]) > 100 or any( not isinstance(project, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,127}", project) for project in scope["projects"] ) or len(set(scope["projects"])) != len(scope["projects"]) or type(scope["include_shared"]) is not bool ): raise ValueError("checkpoint scope is invalid") for field in ("sources", "pending_operations", "proposals", "outputs"): if not isinstance(checkpoint[field], list) or len(checkpoint[field]) > 10_000: raise ValueError(f"{field} must be a bounded list") for source in checkpoint["sources"]: if ( not isinstance(source, dict) or not isinstance(source.get("source_revision_id"), str) or not isinstance(source.get("status"), str) ): raise ValueError("source records require a revision ID and capture status") extraction_id = source.get("extraction_id") if extraction_id is not None and not _SHA256.fullmatch(str(extraction_id)): raise ValueError("source extraction references must use a SHA-256 ID") if not isinstance(checkpoint["coverage"], dict) or set(checkpoint["coverage"]) != { "selected", "included", "excluded", "skipped", }: raise ValueError("coverage must record selected, included, excluded and skipped items") for key, items in checkpoint["coverage"].items(): if not isinstance(items, list) or len(items) > 10_000: raise ValueError("coverage categories must be bounded lists") if key in {"selected", "included"} and any(not isinstance(item, str) for item in items): raise ValueError("selected and included coverage entries must be strings") if key in {"excluded", "skipped"} and any( not isinstance(item, dict) or not isinstance(item.get("item"), str) or not isinstance(item.get("reason"), str) or not item["reason"].strip() for item in items ): raise ValueError("excluded and skipped coverage items require a reason") if not isinstance(checkpoint["budgets"], dict) or len(checkpoint["budgets"]) > 100: raise ValueError("budgets must be a bounded object") if any(type(value) is not int or value < 0 for value in checkpoint["budgets"].values()): raise ValueError("budget values must be non-negative integers") if ( not isinstance(checkpoint["runtime"], dict) or set(checkpoint["runtime"]) != {"host", "model"} or any( value is not None and not isinstance(value, str) for value in checkpoint["runtime"].values() ) ): raise ValueError( "runtime metadata must identify the host and model or explicitly omit them" ) if not isinstance(checkpoint["stage"], str) or not checkpoint["stage"]: raise ValueError("stage must be explicit") for operation in checkpoint["pending_operations"]: if ( not isinstance(operation, dict) or not isinstance(operation.get("operation_id"), str) or not isinstance(operation.get("tool_name"), str) or not _SHA256.fullmatch(str(operation.get("input_sha256", ""))) or not isinstance(operation.get("status"), str) or operation["status"] not in {"pending", "uncertain", "complete", "failed"} ): raise ValueError("pending operation identity, input hash and status are required") for proposal in checkpoint["proposals"]: if ( not isinstance(proposal, dict) or not isinstance(proposal.get("proposal_id"), str) or type(proposal.get("revision")) is not int or proposal["revision"] < 1 or not isinstance(proposal.get("status"), str) or proposal["status"] not in { "staged", "evaluated", "eligible", "needs-review", "invalid", "conflict", "publishing", "accepted", "rejected", "withdrawn", "uncertain", } or ( proposal.get("status") == "accepted" and not isinstance(proposal.get("receipt"), dict) ) ): raise ValueError( "proposal records require an ID, revision, status and accepted receipt" ) if proposal["status"] == "accepted": receipt = proposal["receipt"] if ( receipt.get("proposal_id") != proposal["proposal_id"] or receipt.get("revision") != proposal["revision"] or receipt.get("state") != "accepted" or not re.fullmatch( r"(?:[a-f0-9]{40}|[a-f0-9]{64})", str(receipt.get("commit_id", "")) ) ): raise ValueError("accepted proposal receipt does not match its revision") for output in checkpoint["outputs"]: if not isinstance(output, dict) or not _SHA256.fullmatch(str(output.get("sha256", ""))): raise ValueError("output records require a SHA-256") def main() -> None: parser = argparse.ArgumentParser(description=__doc__) subparsers = parser.add_subparsers(dest="command", required=True) init = subparsers.add_parser("init") init.add_argument("--workspace", type=Path, default=Path.cwd()) init.add_argument("--run-id", required=True) init.add_argument("--objective", required=True) init.add_argument("--projects", required=True) init.add_argument("--include-shared", action="store_true") init.add_argument("--lens", required=True) init.add_argument("--budgets-json", default="{}") update = subparsers.add_parser("update") update.add_argument("--workspace", type=Path, default=Path.cwd()) update.add_argument("--run-id", required=True) show = subparsers.add_parser("show") show.add_argument("--workspace", type=Path, default=Path.cwd()) show.add_argument("--run-id", required=True) arguments = parser.parse_args() if arguments.command == "init": checkpoint = create_checkpoint( arguments.workspace, run_id=arguments.run_id, objective=arguments.objective, projects=[item for item in arguments.projects.split(",") if item], include_shared=arguments.include_shared, lens=arguments.lens, budgets=json.loads(arguments.budgets_json), ) elif arguments.command == "show": checkpoint = load_checkpoint(arguments.workspace, arguments.run_id) else: patch_text = sys.stdin.read(64_001) if len(patch_text.encode("utf-8")) > 64_000: raise SystemExit("checkpoint update exceeds 64 KiB") patch = json.loads(patch_text, object_pairs_hook=_reject_duplicate_keys) if not isinstance(patch, dict): raise SystemExit("checkpoint update must be one JSON object on stdin") checkpoint = update_checkpoint(arguments.workspace, arguments.run_id, patch) print(json.dumps(checkpoint, ensure_ascii=False, sort_keys=True)) if __name__ == "__main__": main()