diff --git a/chronicle/__init__.py b/chronicle/__init__.py index f44b2fd..c4125ac 100644 --- a/chronicle/__init__.py +++ b/chronicle/__init__.py @@ -22,6 +22,7 @@ from chronicle.envelope.store import EnvelopeStore from chronicle.execution_graph import ExecutionGraph from chronicle.redaction import apply_redactors, default_redactors, redact_secrets +from chronicle.replay.checksum import ChecksumMismatch from chronicle.replay.plan import BoundaryMode, ReplayPlan from chronicle.session import ChronicleSession, SessionMode, get_session, reset_session from chronicle.wrap import instrument_langgraph, wrap @@ -31,6 +32,7 @@ __all__ = [ "ActionResult", "BoundaryMode", + "ChecksumMismatch", "ChronicleSession", "ContextMetadata", "Envelope", diff --git a/chronicle/envelope/canonical.py b/chronicle/envelope/canonical.py new file mode 100644 index 0000000..6423bf0 --- /dev/null +++ b/chronicle/envelope/canonical.py @@ -0,0 +1,47 @@ +"""Canonical, volatile-insensitive hashing of recorded inputs. + +The replay checksum hashes each crossing's recorded input. A faithful replay must not +raise merely because a volatile field (a timestamp, a generated identifier, an absolute +path) differs, so such fields are dropped before hashing. Dictionaries are key-sorted so +that field order does not affect the digest. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Iterable +from typing import Any + +# Field names dropped before hashing. Extend via canonicalize(..., volatile=...). +DEFAULT_VOLATILE: frozenset[str] = frozenset( + { + "timestamp", + "created_at", + "updated_at", + "time", + "date", + "request_id", + "span_id", + "run_id", + "uuid", + "nonce", + } +) + + +def canonicalize(value: Any, volatile: Iterable[str] = DEFAULT_VOLATILE) -> Any: + """Return a copy of ``value`` with volatile dict keys dropped, recursively.""" + volatile = set(volatile) + if isinstance(value, dict): + return {k: canonicalize(v, volatile) for k, v in value.items() if k not in volatile} + if isinstance(value, (list, tuple)): + return [canonicalize(v, volatile) for v in value] + return value + + +def digest(value: Any, volatile: Iterable[str] = DEFAULT_VOLATILE) -> str: + """SHA-256 of the canonicalized value, with dict keys sorted for stability.""" + canonical = canonicalize(value, volatile) + encoded = json.dumps(canonical, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() diff --git a/chronicle/execution_graph.py b/chronicle/execution_graph.py index 92a5d63..930a8ce 100644 --- a/chronicle/execution_graph.py +++ b/chronicle/execution_graph.py @@ -67,6 +67,17 @@ def load(cls, directory: str | Path) -> ExecutionGraph: for eid, node in graph.nodes.items() if node.envelope.parent_envelope_id is None ] + + stored = meta.get("checksum") + if stored is not None: + from chronicle.replay.checksum import ChecksumMismatch, trace_checksum + + recomputed = trace_checksum([node.envelope for node in graph.nodes.values()]) + if recomputed != stored: + raise ChecksumMismatch( + f"fixture checksum mismatch in {root}: a recorded input was modified " + "after the trace was committed" + ) return graph def save(self, directory: str | Path) -> None: @@ -97,11 +108,14 @@ def save(self, directory: str | Path) -> None: if node.envelope.parent_envelope_id: edges.append([node.envelope.parent_envelope_id, node.envelope.envelope_id]) + from chronicle.replay.checksum import trace_checksum + graph_json = { "trace_id": self.trace_id, "nodes": node_entries, "edges": edges, "roots": self.root_ids, + "checksum": trace_checksum([node.envelope for node in ordered]), } (root / "graph.json").write_text(json.dumps(graph_json, indent=2)) diff --git a/chronicle/replay/checksum.py b/chronicle/replay/checksum.py new file mode 100644 index 0000000..b2c8b48 --- /dev/null +++ b/chronicle/replay/checksum.py @@ -0,0 +1,86 @@ +"""Order-sensitive checksum over the crossings a replay stubs. + +Two guards use it. (1) Fixture integrity: at record time a digest of the ordered +crossings is stored with the trace; loading a trace recomputes it and raises if a +committed fixture was edited in a way that changes a recorded input (volatile fields are +normalized first, so a timestamp edit does not false-fire). (2) Replay order: during +replay, ``ReplayVerifier`` checks that the crossings actually stubbed, in request order, +match the recorded subsequence the plan stubs, so an insertion, removal, or reordering +raises before a stub returns a value recorded for a different crossing. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Callable +from dataclasses import dataclass + +from chronicle.envelope.canonical import digest +from chronicle.envelope.schema import Envelope + + +class ChecksumMismatch(RuntimeError): + """Replay or a fixture diverged from the recorded crossings.""" + + +def crossing_key(envelope: Envelope) -> str: + """A stable identity for one recorded crossing: boundary, index, canonical input.""" + input_data = envelope.input_state.model_dump() + return f"{envelope.node_id}#{envelope.invocation_index}:{digest(input_data)}" + + +def trace_checksum(envelopes: list[Envelope]) -> str: + """Order-sensitive digest over the crossings in recorded (sequence) order.""" + hasher = hashlib.sha256() + for envelope in sorted(envelopes, key=lambda e: e.sequence): + hasher.update(crossing_key(envelope).encode("utf-8")) + hasher.update(b"|") + return hasher.hexdigest() + + +@dataclass(frozen=True) +class _Expected: + boundary_id: str + invocation_index: int + key: str + + +class ReplayVerifier: + """Checks stubbed crossings against the recorded stubbed subsequence, in order.""" + + def __init__(self, expected: list[_Expected]) -> None: + self._expected = expected + self._pos = 0 + + def check(self, boundary_id: str, invocation_index: int, envelope: Envelope) -> None: + if self._pos >= len(self._expected): + raise ChecksumMismatch( + f"replay stubbed an extra crossing {boundary_id}@{invocation_index} " + "beyond the recorded stubbed subsequence" + ) + expected = self._expected[self._pos] + if (boundary_id, invocation_index) != (expected.boundary_id, expected.invocation_index): + raise ChecksumMismatch( + f"stubbed crossing out of order: replay reached {boundary_id}@" + f"{invocation_index} where the record has {expected.boundary_id}@" + f"{expected.invocation_index}" + ) + if crossing_key(envelope) != expected.key: + raise ChecksumMismatch( + f"stubbed crossing {boundary_id}@{invocation_index} does not match " + "its recorded input" + ) + self._pos += 1 + + +def build_verifier( + envelopes: list[Envelope], + should_stub: Callable[[str, int], bool], +) -> ReplayVerifier: + """The expected stubbed subsequence: recorded crossings the plan stubs, in order.""" + expected = [ + _Expected(env.node_id, env.invocation_index, crossing_key(env)) + for env in sorted(envelopes, key=lambda e: e.sequence) + if should_stub(env.node_id, env.invocation_index) + ] + return ReplayVerifier(expected) diff --git a/chronicle/session.py b/chronicle/session.py index 3210a6b..13df6ab 100644 --- a/chronicle/session.py +++ b/chronicle/session.py @@ -67,6 +67,9 @@ class ChronicleSession: # reach a committed fixture. Empty by default; set to default_redactors() or # your own. Signature: (str) -> str. See chronicle.redaction. redactors: list[Callable[[str], str]] = field(default_factory=list) + # Guard replay fidelity: verify that stubbed crossings occur in the recorded + # order, so a control-flow divergence cannot pair a stub with the wrong envelope. + verify_checksum: bool = True _sequence: int = 0 _invocation_counts: dict[str, int] = field(default_factory=dict) @@ -76,6 +79,7 @@ class ChronicleSession: _captured_results: dict[tuple[str, int], Any] = field(default_factory=dict) _recorded_envelopes: list[Envelope] = field(default_factory=list) _last_envelope_id: str | None = None + _verifier: Any = None def begin_trace(self, trace_id: str | None = None) -> str: if trace_id: @@ -97,6 +101,7 @@ def enable_replay(self, plan: ReplayPlan | None = None) -> None: self.mode = SessionMode.REPLAY self.replay_plan = plan or ReplayPlan() self._replay_cursor.clear() + self._verifier = None def enable_live(self) -> None: self.mode = SessionMode.LIVE @@ -108,6 +113,7 @@ def load_trace(self, path: str | Path) -> ExecutionGraph: self.fixture_graph = ExecutionGraph.load(path) self.trace_id = self.fixture_graph.trace_id self._replay_cursor.clear() + self._verifier = None return self.fixture_graph def current_parent_id(self) -> str | None: @@ -196,11 +202,22 @@ def _fixture_for(self, boundary_id: str) -> Envelope: cursor = self._replay_cursor.get(boundary_id, 0) + 1 self._replay_cursor[boundary_id] = cursor envelope = self.fixture_graph.envelope(boundary_id, cursor) + if self.verify_checksum: + self._ensure_verifier().check(boundary_id, cursor, envelope) self._call_log.append( CallRecord(boundary_id, cursor, "stub", envelope.envelope_id) ) return envelope + def _ensure_verifier(self): + if self._verifier is None: + from chronicle.replay.checksum import build_verifier + + self._verifier = build_verifier( + self.fixture_graph.timeline(), self.replay_plan.should_stub + ) + return self._verifier + def stub_result(self, boundary_id: str, kind: str) -> Any: envelope = self._fixture_for(boundary_id) return envelope_to_return_value(envelope, kind) diff --git a/tests/test_checksum.py b/tests/test_checksum.py new file mode 100644 index 0000000..6db215b --- /dev/null +++ b/tests/test_checksum.py @@ -0,0 +1,138 @@ +"""Replay checksum: canonicalization ignores volatile fields, a tampered fixture is +detected on load, and stubbed crossings are verified against the recorded order. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +import chronicle +from chronicle import ChecksumMismatch, boundary +from chronicle.envelope.canonical import canonicalize, digest +from chronicle.envelope.schema import InputState +from chronicle.execution_graph import ExecutionGraph +from chronicle.replay.checksum import build_verifier, trace_checksum + + +# --- canonicalization -------------------------------------------------------------- # +def test_canonicalize_drops_volatile_and_sorts(): + a = {"amount": 100, "timestamp": "2026-01-01", "nested": {"run_id": "x", "keep": 1}} + assert canonicalize(a) == {"amount": 100, "nested": {"keep": 1}} + + +def test_digest_is_stable_across_volatile_changes_but_not_real_changes(): + base = {"amount": 100, "timestamp": "t1"} + same_but_volatile = {"amount": 100, "timestamp": "t2"} + real_change = {"amount": 101, "timestamp": "t1"} + assert digest(base) == digest(same_but_volatile) # timestamp ignored + assert digest(base) != digest(real_change) # amount matters + + +# --- fixture integrity on load ----------------------------------------------------- # +def _record(tmp_path: Path, ts: str = "2026-01-01T00:00:00Z") -> str: + @boundary("agent", kind="llm") + def agent(state): + return {**state, "completion": "x", "finish_reason": "stop"} + + @boundary( + "tool", + kind="tool", + extract_input=lambda *a, **k: InputState( + messages=[], graph_state={"amount": 100, "timestamp": ts} + ), + ) + def tool(amount): + return {"ok": amount} + + fixture = str(tmp_path / "trace") + with chronicle.record("t", export=fixture): + state = agent({"messages": []}) + tool(100) + agent(state) + return fixture + + +def _tool_fixture_file(fixture: str) -> Path: + meta = json.loads((Path(fixture) / "graph.json").read_text()) + entry = next(n for n in meta["nodes"] if n["boundary_id"] == "tool") + return Path(fixture) / entry["fixture"] + + +def _edit_tool_input(fixture: str, key: str, value) -> None: + path = _tool_fixture_file(fixture) + data = json.loads(path.read_text()) + data["input_state"]["graph_state"][key] = value + path.write_text(json.dumps(data)) + + +def test_volatile_field_edit_does_not_false_fire(tmp_path): + fixture = _record(tmp_path) + _edit_tool_input(fixture, "timestamp", "totally-different-time") + # Reloading must not raise: the volatile field is normalized before hashing. + ExecutionGraph.load(fixture) + + +def test_nonvolatile_edit_is_detected(tmp_path): + fixture = _record(tmp_path) + _edit_tool_input(fixture, "amount", 999) + with pytest.raises(ChecksumMismatch): + ExecutionGraph.load(fixture) + + +# --- replay-order verifier --------------------------------------------------------- # +def test_verifier_accepts_recorded_order(tmp_path): + fixture = _record(tmp_path) + graph = ExecutionGraph.load(fixture) + verifier = build_verifier(graph.timeline(), lambda _b, _i: True) # stub all + for env in graph.timeline(): + verifier.check(env.node_id, env.invocation_index, env) # in order: no raise + + +def test_verifier_flags_out_of_order(tmp_path): + fixture = _record(tmp_path) + graph = ExecutionGraph.load(fixture) + order = graph.timeline() # agent@1, tool@1, agent@2 + verifier = build_verifier(order, lambda _b, _i: True) + # Serve the second crossing first: a reorder / removal of the first. + with pytest.raises(ChecksumMismatch): + verifier.check(order[1].node_id, order[1].invocation_index, order[1]) + + +def test_verifier_flags_extra_stub(tmp_path): + fixture = _record(tmp_path) + graph = ExecutionGraph.load(fixture) + order = graph.timeline() + verifier = build_verifier(order, lambda _b, _i: True) + for env in order: + verifier.check(env.node_id, env.invocation_index, env) + with pytest.raises(ChecksumMismatch): # one crossing beyond the recorded subsequence + verifier.check(order[0].node_id, 99, order[0]) + + +def test_full_replay_of_a_clean_trace_passes(tmp_path): + fixture = _record(tmp_path) + with chronicle.replay_trace(fixture, chronicle.ReplayPlan()): + # Re-run the same shape; every stubbed crossing matches the recorded order. + @boundary("agent", kind="llm") + def agent(state): + return {**state} + + @boundary("tool", kind="tool") + def tool(amount): + return {"ok": amount} + + state = agent({"messages": []}) + tool(100) + agent(state) # no ChecksumMismatch + + +def test_checksum_present_in_graph_json(tmp_path): + fixture = _record(tmp_path) + meta = json.loads((Path(fixture) / "graph.json").read_text()) + assert isinstance(meta.get("checksum"), str) and len(meta["checksum"]) == 64 + # Recomputing from the committed envelopes reproduces it. + graph = ExecutionGraph.load(fixture) + assert trace_checksum(graph.timeline()) == meta["checksum"]