Blame for tiny_computer_agent/audit.py

"""A run's record, chained so that editing it is louder than deleting it.

The harness already refuses to take a brain's word for success, and a trace already keeps a
run's decisions so they can be replayed and scored. What neither does is make the *record*
checkable. A JSON file can be edited by whoever would prefer a better result, and nothing
notices; a directory of screenshots can lose the one frame that proves a failure, and nothing
notices either.

An audit log is an append-only JSONL whose every line carries the digest of the line before
it. Change one entry and its own digest stops matching; delete one and the next entry's `prev`
points at an entry that is not there; reorder them and the indices give it away.
:func:`verify_audit` says which entry is at fault, by number, because "the log is fine" is
only useful if a reader can check it.

    log = AuditLog(Path("runs/today.jsonl"))
    loop = AgentLoop(screen=..., brain=..., audit=log)     # one entry per step
    loop.run("click 'Save' to save the document")

    ok, detail = verify_audit("runs/today.jsonl")          # (True, "3 entries, chain intact")

What this is not: a signature. A hash chain catches edits, not an author who simply writes a
whole new chain. It is evidence against accident and casual tampering, and it says so.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

from tiny_computer_agent.observe import Observation

__all__ = ["AuditEntry", "AuditLog", "observation_digest", "read_audit", "verify_audit"]


def _canonical(payload: dict[str, Any]) -> str:
    """One stable spelling per entry, so a digest means the same thing twice."""
    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def observation_digest(observation: Observation) -> str:
    """A digest of what was on screen: the element index, not the pixels.

    Two frames with the same labels in the same places hash the same, which is the point -
    a step is tied to the screen it was decided on without storing a screenshot.
    """
    index = [[element.label, element.kind, element.x, element.y, element.width, element.height]
             for element in observation.elements]
    return hashlib.sha256(_canonical({"index": index, "text": observation.text[:200]})
                          .encode("utf-8")).hexdigest()[:16]


@dataclass(frozen=True)
class AuditEntry:
    """One line of the log: a step, or the episode that closed it."""

    kind: str                       # "step" | "episode"
    index: int
    at: str
    goal: str
    screen: str
    action: str
    outcome: str
    ok: bool
    prev: str
    digest: str = ""

    def payload(self) -> dict[str, Any]:
        """Everything the digest covers - which is everything except the digest."""
        return {"kind": self.kind, "index": self.index, "at": self.at, "goal": self.goal,
                "screen": self.screen, "action": self.action, "outcome": self.outcome,
                "ok": self.ok, "prev": self.prev}

    def compute_digest(self) -> str:
        return hashlib.sha256(_canonical(self.payload()).encode("utf-8")).hexdigest()

    def to_dict(self) -> dict[str, Any]:
        return {**self.payload(), "digest": self.digest}

    @classmethod
    def from_dict(cls, raw: dict[str, Any]) -> "AuditEntry":
        return cls(kind=raw.get("kind", "step"), index=raw.get("index", -1),
                   at=raw.get("at", ""), goal=raw.get("goal", ""), screen=raw.get("screen", ""),
                   action=raw.get("action", ""), outcome=raw.get("outcome", ""),
                   ok=bool(raw.get("ok", False)), prev=raw.get("prev", ""),
                   digest=raw.get("digest", ""))


@dataclass
class AuditLog:
    """An append-only, hash-chained record of what an episode did."""

    path: Path | str
    _last: str = field(default="", init=False, repr=False)
    _count: int = field(default=0, init=False, repr=False)

    def __post_init__(self) -> None:
        self.path = Path(self.path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        # resuming a log keeps its chain: pick up the last digest rather than starting over
        if self.path.exists():
            existing = read_audit(self.path)
            self._count = len(existing)
            if existing:
                self._last = existing[-1].digest

    # -- writing ----------------------------------------------------------- #

    def _write(self, kind: str, *, goal: str, screen: str = "", action: str = "",
               outcome: str = "", ok: bool = False) -> AuditEntry:
        entry = AuditEntry(kind=kind, index=self._count,
                           at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
                           goal=goal, screen=screen, action=action, outcome=outcome,
                           ok=ok, prev=self._last)
        entry = AuditEntry(**{**entry.payload(), "digest": entry.compute_digest()})
        with open(self.path, "a", encoding="utf-8") as handle:
            handle.write(_canonical(entry.to_dict()) + "\n")
        self._last, self._count = entry.digest, self._count + 1
        return entry

    def append(self, step: Any, goal: str = "") -> AuditEntry:
        """Record one step. Accepts the loop's :class:`~tiny_computer_agent.loop.Step`."""
        result = getattr(step, "result", None)
        outcome = (getattr(step, "error", None) or (result.summary() if result else "")) or ""
        return self._write("step", goal=goal or "",
                           screen=observation_digest(step.observation),
                           action=step.action.describe() if step.action else "",
                           outcome=str(outcome), ok=bool(step.ok))

    def close(self, episode: Any) -> AuditEntry:
        """Record how the episode ended - verified, failed, refused, or out of budget."""
        return self._write("episode", goal=getattr(episode, "goal", ""),
                           outcome=f"{'success' if episode.success else 'failure'}: "
                                   f"{episode.stop_reason}",
                           ok=bool(episode.success))

    # -- reading ----------------------------------------------------------- #

    def entries(self) -> list[AuditEntry]:
        return read_audit(self.path)

    def verify(self) -> tuple[bool, str]:
        return verify_audit(self.path)

    def render(self) -> str:
        entries = self.entries()
        steps = [e for e in entries if e.kind == "step"]
        ok, detail = self.verify()
        counts: dict[str, int] = {}
        for entry in steps:
            label = "ok" if entry.ok else "failed"
            counts[label] = counts.get(label, 0) + 1
        tally = ", ".join(f"{n} {label}" for label, n in sorted(counts.items())) or "no steps"
        return (f"audit: {len(entries)} entries ({len(steps)} steps, {tally}) | "
                f"chain {'intact' if ok else 'BROKEN'}: {detail}")


def read_audit(path: Path | str) -> list[AuditEntry]:
    """Every entry in the log, in order. A malformed line is an entry that fails to load."""
    path = Path(path)
    if not path.exists():
        return []
    entries: list[AuditEntry] = []
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                entries.append(AuditEntry.from_dict(json.loads(line)))
    return entries


def verify_audit(path: Path | str) -> tuple[bool, str]:
    """Check the chain. Returns ``(ok, detail)``, naming the first entry at fault.

    Deliberately reports *which* entry broke and *how* - a digest that does not match, a
    `prev` pointing at nothing, or an index out of order - because those three failures mean
    different things to whoever has to decide whether to trust the run.
    """
    path = Path(path)
    if not path.exists():
        return False, f"no audit log at {path}"

    through = ""
    count = 0
    for number, line in enumerate((l for l in path.read_text(encoding="utf-8").splitlines()
                                   if l.strip()), start=1):
        entry = AuditEntry.from_dict(json.loads(line))
        if entry.index != number - 1:
            return False, (f"entry {number} is out of order (index {entry.index} where "
                           f"{number - 1} was expected): lines were removed or reordered")
        if entry.prev != through:
            return False, (f"entry {number} does not follow entry {number - 1}: its link "
                           f"points at an entry that is not there")
        if entry.compute_digest() != entry.digest:
            return False, f"entry {number} was edited after it was written"
        through = entry.digest
        count += 1
    if count == 0:
        return False, f"{path} is empty"
    return True, f"{count} entries, chain intact"