atlas-iac/services/hermes/scripts/scm_resume_bootstrap.py

#!/usr/bin/env python3
"""Recover one completed SCM receipt from a mediator-owned worker result log."""

from __future__ import annotations

import argparse
import json
import os
import re
import stat
import sys
from pathlib import Path
from typing import Any

from execution_pool_protocol import ProtocolError, canonical_json, read_key, truncate_utf8
from execution_pool_scm import Boundary

ROOT = Path(os.environ.get("HERMES_WORKER_ROOT", "/workspace"))
KEY_PATH = Path(os.environ.get("HERMES_EXECUTION_POOL_KEY_FILE", "/pool-access/execution-pool-key"))
MAX_LOG_BYTES = 4 * 1024 * 1024
PART = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}\Z")
RESULT_FIELDS = frozenset({"status", "summary", "changed_files", "tests_run", "artifacts", "findings", "blockers"})


def _result_candidate(value: Any) -> dict[str, Any] | None:
    """Accept a direct result or the one canonical Claude session result field."""
    if not isinstance(value, dict):
        return None
    candidate = value.get("structured_output") if "structured_output" in value else value
    if not isinstance(candidate, dict) or set(candidate) != RESULT_FIELDS:
        return None
    if candidate.get("status") != "completed":
        return None
    if not isinstance(candidate.get("summary"), str) or not candidate["summary"].strip():
        return None
    if any(
        not isinstance(candidate.get(name), list)
        or any(not isinstance(item, str) for item in candidate[name])
        for name in RESULT_FIELDS - {"status", "summary"}
    ):
        return None
    return candidate


def _completed_result(assignment: dict[str, Any]) -> dict[str, Any]:
    """Return exactly one completed result JSON object from an ordinal's safe log."""
    parts = tuple(str(assignment[name]) for name in ("board", "task_id", "run_id"))
    if any(not PART.fullmatch(part) for part in parts):
        raise ProtocolError("bootstrap assignment path binding is invalid")
    if ROOT.is_symlink():
        raise ProtocolError("bootstrap workspace root is unsafe")
    root = ROOT.resolve()
    state_root = ROOT / "session-state"
    board_root = state_root / parts[0]
    task_root = board_root / parts[1]
    log_path = task_root / f"{parts[2]}.log"
    if any(path.is_symlink() for path in (state_root, board_root, task_root, log_path)):
        raise ProtocolError("bootstrap result log is unsafe")
    try:
        log_path.resolve(strict=False).relative_to(root)
    except ValueError as error:
        raise ProtocolError("bootstrap result log escapes workspace root") from error
    descriptor = os.open(log_path, os.O_RDONLY | os.O_NONBLOCK | getattr(os, "O_NOFOLLOW", 0))
    try:
        info = os.fstat(descriptor)
        if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_LOG_BYTES:
            raise ProtocolError("bootstrap result log is invalid")
        raw = os.read(descriptor, MAX_LOG_BYTES + 1)
    finally:
        os.close(descriptor)
    candidates: dict[bytes, dict[str, Any]] = {}
    for line in raw.splitlines():
        try:
            value = json.loads(line)
        except (UnicodeDecodeError, json.JSONDecodeError):
            continue
        candidate = _result_candidate(value)
        if candidate is None:
            continue
        candidates[canonical_json(candidate)] = candidate
    if len(candidates) != 1:
        raise ProtocolError("bootstrap result evidence is absent or ambiguous")
    return next(iter(candidates.values()))


def bootstrap(key: bytes, raw_assignment: Any) -> dict[str, Any]:
    """Verify assignment HMAC, derive evidence locally, and emit only its receipt."""
    boundary = Boundary(key)
    assignment = boundary.verify(raw_assignment)
    structured = _completed_result(assignment)
    title = truncate_utf8(structured["summary"].strip()[:240], 507)
    body = json.dumps(structured, indent=2, sort_keys=True)
    return boundary.resume_artifact(assignment, {"title": title, "body": body}, structured)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--assignment-stdin", action="store_true", required=True)
    args = parser.parse_args()
    raw = sys.stdin.buffer.read(64 * 1024 + 1)
    if not args.assignment_stdin or len(raw) > 64 * 1024:
        raise SystemExit("signed assignment input is invalid")
    try:
        receipt = bootstrap(read_key(KEY_PATH), json.loads(raw))
    except (OSError, ProtocolError, ValueError, json.JSONDecodeError) as error:
        raise SystemExit(f"bootstrap rejected: {error}") from error
    sys.stdout.buffer.write(canonical_json(receipt) + b"\n")
    return 0


if __name__ == "__main__":  # pragma: no cover - operator entrypoint
    raise SystemExit(main())