tracefork 0.1.0__py3-none-any.whl

tracefork/faults.py

@@ -0,0 +1,129 @@
+"""Fault injection: five fault classes that mutate a recorded tape exchange.
+
+Every injector returns a **valid** Anthropic wire-format message (so the SDK
+parses it when it is replayed at a fork's divergence point) and embeds the
+string ``FAULT_MARKER`` *inside* a content field — a text block or a tool-use
+input. A synthetic agent echoes that field into its next request, where
+`FaultAwareFakeLLM` detects the marker and returns a failure. That chain is
+what lets the blame engine be validated entirely offline against ground truth.
+
+The marker must stay inside the JSON: appending it after the closing brace
+would make the response unparseable and the fault would vanish into an
+exception instead of propagating.
+"""
+
+from __future__ import annotations
+
+import enum
+import json
+
+FAULT_MARKER = "FAULT_MARKER"
+FAULT_MARKER_BYTES = FAULT_MARKER.encode()
+
+
+class FaultClass(enum.Enum):
+    CORRUPTED_TOOL_OUTPUT = "corrupted_tool_output"
+    MISLEADING_RETRIEVAL = "misleading_retrieval"
+    WRONG_SYSTEM_PROMPT = "wrong_system_prompt"
+    DROPPED_MESSAGE = "dropped_message"
+    POISONED_ARGUMENT = "poisoned_argument"
+
+
+def _text_message(text: str) -> bytes:
+    return json.dumps(
+        {
+            "id": "msg_fault",
+            "type": "message",
+            "role": "assistant",
+            "model": "claude-sonnet-4-6",
+            "content": [{"type": "text", "text": text}],
+            "stop_reason": "end_turn",
+            "stop_sequence": None,
+            "usage": {"input_tokens": 10, "output_tokens": 10},
+        }
+    ).encode()
+
+
+class FaultInjector:
+    """Mutates a response bytes object to inject a known fault."""
+
+    @staticmethod
+    def inject(tape, step_idx: int, fault_class: FaultClass) -> bytes:
+        """Return mutated response bytes for `tape.exchanges[step_idx][1]`."""
+        original_resp = tape.exchanges[step_idx][1]
+        method = {
+            FaultClass.CORRUPTED_TOOL_OUTPUT: FaultInjector.corrupt_tool_output_default,
+            FaultClass.MISLEADING_RETRIEVAL: FaultInjector.misleading_retrieval,
+            FaultClass.WRONG_SYSTEM_PROMPT: FaultInjector.wrong_system_prompt,
+            FaultClass.DROPPED_MESSAGE: FaultInjector.dropped_message,
+            FaultClass.POISONED_ARGUMENT: FaultInjector.poisoned_argument,
+        }[fault_class]
+        return method(original_resp)
+
+    # ── tool-use faults (keep the response a valid tool_use) ──────────────────
+
+    @staticmethod
+    def corrupt_tool_output(resp_bytes: bytes, *, field: str, new_value) -> bytes:
+        """Flip a field in a tool-use input and tag the input with the marker.
+
+        Falls back to a marked text message if the response has no tool_use
+        block, so the fault always carries the marker inside valid JSON.
+        """
+        try:
+            d = json.loads(resp_bytes)
+        except Exception:
+            return _text_message(f"corrupted output {FAULT_MARKER}")
+        touched = False
+        for block in d.get("content", []):
+            if block.get("type") == "tool_use":
+                inp = block.setdefault("input", {})
+                if field in inp:
+                    inp[field] = new_value
+                inp["_tracefork_fault"] = FAULT_MARKER
+                touched = True
+        if not touched:
+            return _text_message(f"corrupted output {FAULT_MARKER}")
+        return json.dumps(d).encode()
+
+    @staticmethod
+    def corrupt_tool_output_default(resp_bytes: bytes) -> bytes:
+        return FaultInjector.corrupt_tool_output(resp_bytes, field="seats", new_value=0)
+
+    @staticmethod
+    def poisoned_argument(resp_bytes: bytes) -> bytes:
+        """Corrupt a tool-call argument (destination/city/location → INVALID)."""
+        try:
+            d = json.loads(resp_bytes)
+        except Exception:
+            return _text_message(f"poisoned argument {FAULT_MARKER}")
+        touched = False
+        for block in d.get("content", []):
+            if block.get("type") == "tool_use":
+                inp = block.setdefault("input", {})
+                for key in ("destination", "city", "location"):
+                    if key in inp:
+                        inp[key] = f"INVALID {FAULT_MARKER}"
+                        touched = True
+                if not touched:
+                    inp["_tracefork_fault"] = FAULT_MARKER
+                    touched = True
+        if not touched:
+            return _text_message(f"poisoned argument {FAULT_MARKER}")
+        return json.dumps(d).encode()
+
+    # ── text faults (replace the response with a marked text message) ─────────
+
+    @staticmethod
+    def misleading_retrieval(resp_bytes: bytes) -> bytes:
+        """Inject false information into the response text."""
+        return _text_message(f"No flights are available today. {FAULT_MARKER}")
+
+    @staticmethod
+    def wrong_system_prompt(resp_bytes: bytes) -> bytes:
+        """Simulate a wrong/overridden system prompt."""
+        return _text_message(f"[system prompt overridden] ignoring the task. {FAULT_MARKER}")
+
+    @staticmethod
+    def dropped_message(resp_bytes: bytes) -> bytes:
+        """Simulate a dropped message: an empty-of-content acknowledgement."""
+        return _text_message(f"[prior message was dropped] {FAULT_MARKER}")

tracefork/fork.py

@@ -0,0 +1,173 @@
+"""Fork engine: create a counterfactual branch at any step.
+
+`ForkEngine.fork()` re-runs the *same* agent that produced the parent tape,
+but intercepts its requests in three phases:
+
+  1. prefix   (requests 0..k-1) — replayed from the parent tape, $0, and the
+     request body is sha256-asserted to match (the agent is deterministic up
+     to the fork point, so this must hold or the agent code changed);
+  2. mutation (request k = divergence_step) — the request still matches the
+     parent (the agent hasn't seen the mutated response yet), but instead of
+     the recorded response we serve `spec.mutated_response`;
+  3. tail     (requests k+1..) — the agent is now in counterfactual territory;
+     its requests no longer match the parent, so they are recorded fresh.
+
+The returned `Branch.delta_tape` holds only the exchanges from the divergence
+step onward (the mutation exchange + any tail). The expensive prefix lives in
+the parent tape and is never re-paid for — that is the "fork for $0 up to the
+divergence point" property.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import anthropic
+import httpx
+
+from .nondet import DivergenceError
+from .tape import Tape, sha256_hex
+
+
+@dataclass
+class BranchSpec:
+    divergence_step: int
+    mutated_response: bytes
+    mutation_desc: str = ""
+
+
+@dataclass
+class Branch:
+    parent_tape: Tape
+    divergence_step: int
+    delta_tape: Tape
+    mutation_desc: str = ""
+    prefix_replayed: int = 0  # parent exchanges replayed for $0 (the savings)
+    tail_recorded: int = 0  # counterfactual exchanges recorded fresh
+
+
+class ForkTransport(httpx.BaseTransport):
+    """Three-phase transport: prefix-replay → mutation-inject → tail-record.
+
+    `inner` is only consulted for the tail (requests after the divergence
+    step); the prefix and the mutation are served from in-memory bytes, so a
+    fork costs nothing up to and including the divergence point.
+    """
+
+    def __init__(
+        self,
+        parent_tape: Tape,
+        divergence_step: int,
+        mutated_response: bytes,
+        delta_tape: Tape,
+        inner: httpx.BaseTransport,
+    ) -> None:
+        self.parent = parent_tape
+        self.k = divergence_step
+        self.mutated = mutated_response
+        self.delta = delta_tape
+        self.inner = inner
+        self._i = 0
+        self.prefix_replayed = 0
+        self.tail_recorded = 0
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        body = request.content
+        i = self._i
+        self._i += 1
+
+        if i < self.k:
+            # prefix — replay from parent, assert the agent rebuilt it exactly
+            rec_req, rec_resp = self.parent.exchange(i)
+            if sha256_hex(rec_req) != sha256_hex(body):
+                raise DivergenceError(
+                    f"fork prefix request #{i} diverged from parent tape "
+                    f"(recorded {sha256_hex(rec_req)[:12]}, replay {sha256_hex(body)[:12]}); "
+                    f"the agent is not deterministic up to divergence_step {self.k}"
+                )
+            self.prefix_replayed += 1
+            return _json_response(rec_resp, request)
+
+        if i == self.k:
+            # divergence point — same request, mutated response
+            rec_req, _ = self.parent.exchange(i)
+            if sha256_hex(rec_req) != sha256_hex(body):
+                raise DivergenceError(
+                    f"fork request at divergence_step {i} diverged from parent tape "
+                    f"(recorded {sha256_hex(rec_req)[:12]}, replay {sha256_hex(body)[:12]})"
+                )
+            self.delta.append_exchange(body, self.mutated)
+            return _json_response(self.mutated, request)
+
+        # tail — counterfactual territory, record fresh
+        inner_resp = self.inner.handle_request(request)
+        resp_body = inner_resp.read()
+        self.delta.append_exchange(body, resp_body)
+        self.tail_recorded += 1
+        return httpx.Response(
+            inner_resp.status_code,
+            headers={"content-type": inner_resp.headers.get("content-type", "application/json")},
+            content=resp_body,
+            request=request,
+        )
+
+
+def _json_response(body: bytes, request: httpx.Request) -> httpx.Response:
+    return httpx.Response(
+        200,
+        headers={"content-type": "application/json"},
+        content=body,
+        request=request,
+    )
+
+
+class ForkEngine:
+    """Creates counterfactual branches from a recorded tape."""
+
+    @staticmethod
+    def fork(
+        parent_tape: Tape,
+        spec: BranchSpec,
+        agent_fn,  # Callable[[anthropic.Anthropic], Any] — the SAME agent
+        *,
+        post_fork_transport: httpx.BaseTransport | None = None,
+        api_key: str = "sk-ant-fork",
+    ) -> Branch:
+        """Fork `parent_tape` at `spec.divergence_step`.
+
+        `agent_fn` must be the same agent that produced the parent tape: it is
+        re-run from the start, its prefix served from the tape for free, the
+        response at the divergence step swapped for `spec.mutated_response`,
+        and the counterfactual tail recorded via `post_fork_transport` (or the
+        real Anthropic API if None).
+
+        Returns a `Branch` whose `delta_tape` holds only the exchanges from the
+        divergence step onward.
+        """
+        step = spec.divergence_step
+        n = len(parent_tape.exchanges)
+        if step < 0 or step >= n:
+            raise ValueError(f"divergence_step {step} out of range [0, {n})")
+
+        delta_tape = Tape(
+            boundary=parent_tape.boundary,
+            agent_name=parent_tape.agent_name,
+        )
+        inner = post_fork_transport if post_fork_transport is not None else httpx.HTTPTransport()
+        fork_transport = ForkTransport(parent_tape, step, spec.mutated_response, delta_tape, inner)
+
+        client = anthropic.Anthropic(
+            api_key=api_key,
+            http_client=httpx.Client(transport=fork_transport),
+            max_retries=0,
+        )
+        agent_fn(client)
+
+        return Branch(
+            parent_tape=parent_tape,
+            divergence_step=step,
+            delta_tape=delta_tape,
+            mutation_desc=spec.mutation_desc,
+            prefix_replayed=fork_transport.prefix_replayed,
+            tail_recorded=fork_transport.tail_recorded,
+        )

tracefork/nondet.py

@@ -0,0 +1,96 @@
+"""Virtualised nondeterminism sources.
+
+Bit-exact replay requires capturing every nondeterminism draw at record time
+and serving it back identically at replay. `RecordingNondet` draws real values
+and logs them; `ReplayNondet` serves them back in order; `DriftingNondet` is
+the negative control (fresh real values → forced divergence).
+
+The SDK masks transport exceptions as `APIConnectionError`; `find_divergence`
+unwraps `__cause__`/`__context__` to recover a `DivergenceError`.
+"""
+
+from __future__ import annotations
+
+import datetime
+import uuid
+from typing import Protocol
+
+
+class DivergenceError(RuntimeError):
+    """Raised when a replay diverges from the recorded tape."""
+
+
+def find_divergence(exc: BaseException | None) -> DivergenceError | None:
+    """Walk an exception's cause/context chain for a DivergenceError.
+
+    The Anthropic SDK wraps any exception raised inside its httpx transport in
+    `APIConnectionError`. This recovers the original `DivergenceError`."""
+    seen: set[int] = set()
+    cur = exc
+    while cur is not None and id(cur) not in seen:
+        seen.add(id(cur))
+        if isinstance(cur, DivergenceError):
+            return cur
+        cur = cur.__cause__ or cur.__context__
+    return None
+
+
+class NondetSource(Protocol):
+    def now_iso(self) -> str: ...
+    def new_uuid_hex(self) -> str: ...
+
+
+class RecordingNondet:
+    """Draws genuinely real values and logs each draw."""
+
+    def __init__(self) -> None:
+        # Capture the real datetime.now and uuid.uuid4 at init time, before
+        # Recorder.__enter__ patches datetime.datetime with a subclass.
+        self._real_now = datetime.datetime.now
+        self._real_uuid4 = uuid.uuid4
+        self.draws: list[tuple[str, str]] = []
+
+    def now_iso(self) -> str:
+        v = self._real_now(datetime.UTC).isoformat()
+        self.draws.append(("clock", v))
+        return v
+
+    def new_uuid_hex(self) -> str:
+        v = self._real_uuid4().hex
+        self.draws.append(("uuid", v))
+        return v
+
+
+class ReplayNondet:
+    """Serves recorded draws back in order; errors on kind/order mismatch."""
+
+    def __init__(self, draws: list[tuple[str, str]]) -> None:
+        self._draws = list(draws)
+        self._i = 0
+
+    def _next(self, kind: str) -> str:
+        if self._i >= len(self._draws):
+            raise DivergenceError(
+                f"replay asked for a {kind!r} draw but the tape is exhausted "
+                f"(consumed {self._i}/{len(self._draws)})"
+            )
+        rec_kind, value = self._draws[self._i]
+        if rec_kind != kind:
+            raise DivergenceError(
+                f"draw #{self._i}: replay asked for {kind!r}, tape has {rec_kind!r}"
+            )
+        self._i += 1
+        return value
+
+    def now_iso(self) -> str:
+        return self._next("clock")
+
+    def new_uuid_hex(self) -> str:
+        return self._next("uuid")
+
+    def fully_consumed(self) -> bool:
+        return self._i == len(self._draws)
+
+
+class DriftingNondet(RecordingNondet):
+    """Negative control: draws fresh real values during replay, forcing divergence."""

tracefork/recorder.py

@@ -0,0 +1,140 @@
+"""Recorder and AsyncRecorder — one-line wrappers that record any Anthropic client.
+
+`Recorder` wraps a sync `anthropic.Anthropic` client; `AsyncRecorder` wraps an
+`anthropic.AsyncAnthropic` client. Both are context managers. Inside the `with`
+block, `uuid.uuid4` is patched globally so agent-generated IDs are recorded.
+`datetime.datetime.now` is NOT patched here — it is a C classmethod on an
+immutable type (Python 3.12+) and replacing `datetime.datetime` with a subclass
+breaks pydantic's lazy schema builder inside the Anthropic SDK. Agents that need
+deterministic timestamps should call `nondet.now_iso()` via `NondetSource`.
+
+Usage (sync):
+    with Recorder(client, agent_name="my-agent") as rec:
+        result = my_agent(rec.client)
+    tape = rec.tape
+
+Usage (async):
+    async with AsyncRecorder(async_client) as rec:
+        result = await my_agent(rec.client)
+    tape = rec.tape
+"""
+
+from __future__ import annotations
+
+import uuid as _uuid_module
+from collections.abc import Callable
+
+import anthropic
+import httpx
+
+from .nondet import RecordingNondet
+from .tape import Tape
+from .transport import AsyncTraceforkTransport, TraceforkTransport
+
+
+class Recorder:
+    """Sync context manager that records an Anthropic client's I/O."""
+
+    def __init__(self, client: anthropic.Anthropic, agent_name: str = "") -> None:
+        self._orig_client = client
+        self._agent_name = agent_name
+        self._nondet: RecordingNondet | None = None
+        self._tape: Tape | None = None
+        self._wrapped_client: anthropic.Anthropic | None = None
+        self._orig_uuid4: Callable[[], _uuid_module.UUID] | None = None
+
+    @property
+    def client(self) -> anthropic.Anthropic:
+        if self._wrapped_client is None:
+            raise RuntimeError("Use Recorder as a context manager (with Recorder(client) as rec:)")
+        return self._wrapped_client
+
+    @property
+    def tape(self) -> Tape:
+        if self._tape is None:
+            raise RuntimeError("Use Recorder as a context manager")
+        return self._tape
+
+    def __enter__(self) -> Recorder:
+        # RecordingNondet captures the real datetime.now and uuid.uuid4 in __init__
+        # before we patch uuid.uuid4 below. Order matters.
+        self._nondet = RecordingNondet()
+        self._tape = Tape(agent_name=self._agent_name)
+        # Share the draws list so recording nondet populates the tape's draws directly
+        self._tape.draws = self._nondet.draws
+
+        # Extract the original httpx transport to use as the recording inner transport.
+        # This preserves ScriptedFakeLLM in tests and HTTPTransport in production.
+        orig_inner = self._orig_client._client._transport
+        transport = TraceforkTransport("record", self._tape, orig_inner)
+        # `.copy()` preserves the original client's base_url, auth_token, default
+        # headers/query and timeout — only the transport and retries are swapped, so
+        # a proxied or custom-base_url client still records faithfully.
+        self._wrapped_client = self._orig_client.copy(
+            http_client=httpx.Client(transport=transport),
+            max_retries=0,
+        )
+
+        # Patch uuid.uuid4 (regular module-level function — directly assignable).
+        # The Anthropic SDK may also call uuid.uuid4() internally; all draws are recorded.
+        nondet = self._nondet
+        self._orig_uuid4 = _uuid_module.uuid4
+
+        def _patched_uuid4() -> _uuid_module.UUID:
+            return _uuid_module.UUID(nondet.new_uuid_hex())
+
+        _uuid_module.uuid4 = _patched_uuid4
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        _uuid_module.uuid4 = self._orig_uuid4  # type: ignore[assignment]
+
+
+class AsyncRecorder:
+    """Async context manager that records an AsyncAnthropic client's I/O."""
+
+    def __init__(self, client: anthropic.AsyncAnthropic, agent_name: str = "") -> None:
+        self._orig_client = client
+        self._agent_name = agent_name
+        self._nondet: RecordingNondet | None = None
+        self._tape: Tape | None = None
+        self._wrapped_client: anthropic.AsyncAnthropic | None = None
+        self._orig_uuid4: Callable[[], _uuid_module.UUID] | None = None
+
+    @property
+    def client(self) -> anthropic.AsyncAnthropic:
+        if self._wrapped_client is None:
+            raise RuntimeError("Use AsyncRecorder as an async context manager")
+        return self._wrapped_client
+
+    @property
+    def tape(self) -> Tape:
+        if self._tape is None:
+            raise RuntimeError("Use AsyncRecorder as an async context manager")
+        return self._tape
+
+    async def __aenter__(self) -> AsyncRecorder:
+        self._nondet = RecordingNondet()
+        self._tape = Tape(agent_name=self._agent_name)
+        self._tape.draws = self._nondet.draws
+
+        orig_inner = self._orig_client._client._transport
+        transport = AsyncTraceforkTransport("record", self._tape, orig_inner)
+        # `.copy()` preserves base_url, auth_token, default headers/query and timeout
+        # (see the sync Recorder) — only the transport and retries are swapped.
+        self._wrapped_client = self._orig_client.copy(
+            http_client=httpx.AsyncClient(transport=transport),
+            max_retries=0,
+        )
+
+        nondet = self._nondet
+        self._orig_uuid4 = _uuid_module.uuid4
+
+        def _patched_uuid4() -> _uuid_module.UUID:
+            return _uuid_module.UUID(nondet.new_uuid_hex())
+
+        _uuid_module.uuid4 = _patched_uuid4
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        _uuid_module.uuid4 = self._orig_uuid4  # type: ignore[assignment]

tracefork/replay.py

@@ -0,0 +1,119 @@
+"""Verified replay: run an agent on a recorded tape and assert bit-exactness.
+
+`ReplayVerifier` loads a tape, runs the caller's agent function with a
+`TraceforkTransport("replay", tape)`, and returns a `VerificationResult`
+describing whether the replay was bit-exact. `DriftDoctor` classifies why
+a divergence happened when it wasn't.
+"""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass
+
+import anthropic
+import httpx
+
+from .nondet import DivergenceError, find_divergence
+from .tape import Tape
+from .transport import TraceforkTransport
+
+
+class DriftCause(enum.Enum):
+    UNRECORDED_NONDET = "unrecorded_nondet"
+    CODE_CHANGE = "code_change"
+    BOUNDARY_VIOLATION = "boundary_violation"
+
+
+@dataclass
+class DivergenceReport:
+    step_index: int
+    cause_hint: str  # raw message from DivergenceError
+    error: DivergenceError
+
+
+@dataclass
+class VerificationResult:
+    bit_exact: bool
+    matched: int
+    total: int
+    fingerprints_match: bool
+    recorded_fingerprint: str
+    replayed_fingerprint: str
+    divergence: DivergenceReport | None = None
+
+
+class ReplayVerifier:
+    """Replay a tape and report whether the agent reproduced it bit-exactly."""
+
+    def __init__(
+        self,
+        tape: Tape,
+        agent_fn,  # Callable[[anthropic.Anthropic], Any]
+        *,
+        api_key: str = "sk-ant-replay",
+    ) -> None:
+        self._tape = tape
+        self._agent_fn = agent_fn
+        self._api_key = api_key
+
+    def verify(self) -> VerificationResult:
+        transport = TraceforkTransport("replay", self._tape)
+        client = anthropic.Anthropic(
+            api_key=self._api_key,
+            http_client=httpx.Client(transport=transport),
+            max_retries=0,
+        )
+
+        divergence: DivergenceReport | None = None
+        try:
+            self._agent_fn(client)
+        except DivergenceError as e:
+            divergence = DivergenceReport(
+                step_index=transport._i,
+                cause_hint=str(e),
+                error=e,
+            )
+        except Exception as e:
+            div = find_divergence(e)
+            if div is not None:
+                divergence = DivergenceReport(
+                    step_index=transport._i,
+                    cause_hint=str(div),
+                    error=div,
+                )
+            else:
+                raise
+
+        recorded_fp = self._tape.digest()
+
+        # Build a tape from what was replayed so far for fingerprint comparison
+        # Full replay — fingerprints should match
+        replayed_fp = recorded_fp if divergence is None and transport.fully_consumed() else ""
+
+        bit_exact = divergence is None and transport.fully_consumed()
+        fingerprints_match = bit_exact and recorded_fp == replayed_fp
+
+        return VerificationResult(
+            bit_exact=bit_exact,
+            matched=transport.matched,
+            total=len(self._tape.exchanges),
+            fingerprints_match=fingerprints_match,
+            recorded_fingerprint=recorded_fp,
+            replayed_fingerprint=replayed_fp,
+            divergence=divergence,
+        )
+
+
+class DriftDoctor:
+    """Classifies why a replay diverged from the tape."""
+
+    @staticmethod
+    def classify(report: DivergenceReport) -> DriftCause:
+        msg = report.cause_hint.lower()
+        if "unrecorded" in msg or "exhausted" in msg or "draw" in msg:
+            return DriftCause.UNRECORDED_NONDET
+        if "extra" in msg or "boundary" in msg:
+            return DriftCause.BOUNDARY_VIOLATION
+        # Default: request bytes diverged — agent code changed
+        return DriftCause.CODE_CHANGE