orchestrator/agent-team/agent_team/state_store.py

"""Atomic state-store utilities and integrity checking (design §6.7).

All durable state on the R720 (LangGraph SQLite checkpoint, the
``pending_questions`` ledger, the budget ledger, the Plane-1 rotation/coverage
pointer) is written atomically (write-temp-then-fsync-then-rename) and
integrity-checked on load. "Integrity-checked" is concrete here: a
schema-version match plus a stored content hash. On any mismatch the loader
refuses to proceed silently and raises :class:`IntegrityError` so the
coordinator can park the affected task with an ALARM rather than acting on
corrupt state.

This module is pure stdlib (``os``, ``tempfile``, ``hashlib``, ``pathlib``)
and depends on no other ``agent_team`` module. The leaf builders import these
signatures verbatim, so they are intentionally explicit and final.
"""

from __future__ import annotations

import hashlib
import json
import os
import tempfile
from pathlib import Path

__all__ = [
    "IntegrityError",
    "atomic_write",
    "compute_content_hash",
    "read_checked",
]

# Sidecar files sit next to the protected payload and carry the integrity
# metadata (schema version + content hash). Keeping them separate from the
# payload means the payload bytes round-trip unchanged.
_META_SUFFIX = ".meta.json"

# Algorithm used for the stored content hash. Recorded in the sidecar so a
# future algorithm change stays backward-readable.
_HASH_ALGO = "sha256"


class IntegrityError(Exception):
    """Raised when durable state fails its integrity check on load.

    Signals a schema-version mismatch, a missing/garbled integrity sidecar, or
    a stored-content-hash mismatch (corruption or tampering). Callers treat
    this as "refuse to proceed silently": park the task and ALARM rather than
    restart blindly (§6.7).
    """


def compute_content_hash(data: bytes) -> str:
    """Return the hex content hash for ``data`` (sha256).

    The same routine is used when writing the sidecar and when verifying on
    load, so the two are guaranteed consistent.
    """
    return hashlib.new(_HASH_ALGO, data).hexdigest()


def _meta_path(path: Path) -> Path:
    """Return the sidecar metadata path for a payload ``path``."""
    return path.with_name(path.name + _META_SUFFIX)


def atomic_write(path: Path, data: bytes) -> None:
    """Atomically write ``data`` to ``path`` (write-temp -> fsync -> rename).

    The bytes are written to a temporary file in the same directory, flushed
    and ``fsync``-ed to durable storage, then ``os.replace``-d onto the final
    path. ``os.replace`` is atomic on POSIX within a filesystem, so a reader
    never observes a half-written file and a crash mid-write leaves either the
    old payload or the new one, never a torn one. The containing directory is
    ``fsync``-ed afterward so the rename itself is durable.

    This writes only the payload; integrity metadata is written by callers via
    :func:`write_checked` / read back by :func:`read_checked`. (The sidecar is
    written through this same primitive, so it is equally crash-safe.)
    """
    path = Path(path)
    directory = path.parent
    directory.mkdir(parents=True, exist_ok=True)

    # delete=False so we control the rename; same dir guarantees same fs.
    fd, tmp_name = tempfile.mkstemp(
        prefix=path.name + ".", suffix=".tmp", dir=directory
    )
    tmp_path = Path(tmp_name)
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(tmp_path, path)
    except BaseException:
        # Best-effort cleanup of the temp file on any failure.
        try:
            os.unlink(tmp_path)
        except FileNotFoundError:
            pass
        raise

    # Make the rename itself durable by fsync-ing the directory.
    dir_fd = os.open(directory, os.O_RDONLY)
    try:
        os.fsync(dir_fd)
    except OSError:
        # Some filesystems disallow directory fsync; the rename is still
        # atomic, only its durability across power-loss is weakened.
        pass
    finally:
        os.close(dir_fd)


def write_checked(path: Path, data: bytes, *, schema_version: int) -> None:
    """Atomically write ``data`` plus its integrity sidecar.

    Writes the payload first, then the sidecar carrying ``schema_version`` and
    the content hash. :func:`read_checked` verifies both. Both writes go
    through :func:`atomic_write`, so each is crash-safe; if a crash lands
    between them the sidecar is simply stale/absent and :func:`read_checked`
    fails closed with :class:`IntegrityError`, which is the intended
    refuse-to-proceed behaviour.
    """
    path = Path(path)
    atomic_write(path, data)
    meta = {
        "schema_version": int(schema_version),
        "hash_algo": _HASH_ALGO,
        "content_hash": compute_content_hash(data),
    }
    atomic_write(_meta_path(path), json.dumps(meta, sort_keys=True).encode("utf-8"))


def read_checked(path: Path, *, schema_version: int) -> bytes:
    """Read and integrity-check ``path``, returning its bytes.

    Verifies the integrity sidecar exists, that its recorded
    ``schema_version`` matches the expected ``schema_version``, and that the
    stored content hash matches a freshly computed hash of the payload bytes.
    Any mismatch (missing/garbled sidecar, schema drift, corruption/tampering)
    raises :class:`IntegrityError`.
    """
    path = Path(path)
    try:
        data = path.read_bytes()
    except FileNotFoundError as exc:
        raise IntegrityError(f"state payload missing: {path}") from exc

    meta_path = _meta_path(path)
    try:
        raw_meta = meta_path.read_bytes()
    except FileNotFoundError as exc:
        raise IntegrityError(f"integrity sidecar missing: {meta_path}") from exc

    try:
        meta = json.loads(raw_meta)
    except (ValueError, UnicodeDecodeError) as exc:
        raise IntegrityError(f"integrity sidecar unreadable: {meta_path}") from exc

    stored_version = meta.get("schema_version")
    if stored_version != schema_version:
        raise IntegrityError(
            f"schema-version mismatch for {path}: "
            f"stored={stored_version!r} expected={schema_version!r}"
        )

    stored_hash = meta.get("content_hash")
    actual_hash = compute_content_hash(data)
    if stored_hash != actual_hash:
        raise IntegrityError(
            f"content-hash mismatch for {path}: "
            f"stored={stored_hash!r} actual={actual_hash!r}"
        )

    return data