feat: checkpoint batching and release-gate stories

tests/test_failure_semantics.py

@@ -0,0 +1,611 @@
+"""Bounded failure, cancellation, and restart semantics (DGR-013).
+
+These tests drive the hardened per-session decode stream with the *same*
+pure-numpy KV-cached dense-Llama reference the Hot KV State manager (DGR-007) and
+the continuous-batch scheduler (DGR-012) use, imported from ``test_hot_kv_state``.
+The whole matrix stays deterministic, download-free, GPU-free, and API-credit-free
+while exercising the real KV isolation path (``KvBoundaryAdapter`` +
+``HotKvStateManager``) rather than a mock.
+
+Coverage maps to the story's acceptance criteria:
+
+* deadlines and heartbeat/health loss terminate blocked stream operations,
+* cancellation propagates across every Shard and releases KV + queued buffers,
+* duplicate steps are idempotent; uncertain mutations are never replayed silently,
+* alpha failover restarts from token zero rather than importing unverified KV,
+* worker death / stream reset / malformed bundle / stale epoch / cache miss,
+* billing/work records distinguish completed, cancelled, failed, and unverified.
+"""
+
+from __future__ import annotations
+
+import json
+
+import numpy as np
+import pytest
+
+from meshnet_node.batch_scheduler import (
+    ContinuousBatchScheduler,
+    DoneReason,
+    GenerationRequest,
+    KvBatchEngine,
+    NodeBudget,
+)
+from meshnet_node.boundary_adapter import BoundaryBundle, BoundaryContractError
+from meshnet_node.hot_kv_state import (
+    CacheMiss,
+    CacheMissReason,
+    HotKvStateConfig,
+    HotKvStateManager,
+    KvBoundaryAdapter,
+    StaleRouteEpochError,
+    kv_recipe_for,
+)
+from meshnet_node.failure_semantics import (
+    CancellationToken,
+    DeadlineGuard,
+    FailureKind,
+    HardenedSessionRunner,
+    IdempotencyLedger,
+    OperationCancelled,
+    RestartController,
+    ShardCancellationGroup,
+    StepKey,
+    StreamTerminated,
+    UncertainMutationError,
+    WorkLedger,
+    WorkRecord,
+    WorkStatus,
+    classify_exception,
+    work_status_for,
+)
+
+# Reuse the certified numpy dense-Llama reference and shard from the DGR-007 gate.
+from test_hot_kv_state import _KvDenseLlama, _KvReferenceShard
+
+
+# --------------------------------------------------------------------------- #
+# Helpers.
+# --------------------------------------------------------------------------- #
+
+
+class _FakeClock:
+    def __init__(self) -> None:
+        self.now = 0.0
+
+    def __call__(self) -> float:
+        return self.now
+
+    def advance(self, delta: float) -> None:
+        self.now += delta
+
+
+class _FaultyShard(_KvReferenceShard):
+    """A full-shard reference that raises on the Nth ``run_layers_cached`` call.
+
+    ``run_layers_cached`` is invoked once per stream step, so ``fail_at_call=k``
+    simulates a worker dying at step ``k-1`` (calls are 1-indexed). The call
+    counter persists across attempts, so a restart on a fresh epoch keeps counting
+    and does not re-trip the same fault.
+    """
+
+    def __init__(self, model, start, end, *, fail_at_call=None, error=None):
+        super().__init__(model, start, end)
+        self._fail_at_call = fail_at_call
+        self._error = error or RuntimeError("worker died mid-step")
+        self.calls = 0
+
+    def run_layers_cached(self, hidden, *, positions, past_kv):
+        self.calls += 1
+        if self._fail_at_call is not None and self.calls == self._fail_at_call:
+            raise self._error
+        return super().run_layers_cached(hidden, positions=positions, past_kv=past_kv)
+
+
+def _make_adapter(model=None, *, config=None, shard=None):
+    """A full-shard KV boundary adapter over the deterministic numpy dense-Llama."""
+    model = model or _KvDenseLlama()
+    shard = shard or _KvReferenceShard(model, 0, model.n_layers - 1)
+    manager = HotKvStateManager(kv_recipe_for(shard), config=config)
+    adapter = KvBoundaryAdapter(shard, manager)
+    return adapter
+
+
+def _generation(session_id, prompt, n_new, epoch=0):
+    return GenerationRequest(
+        session_id=session_id,
+        route_epoch=epoch,
+        prompt_token_ids=tuple(prompt),
+        max_new_tokens=n_new,
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Happy path (the baseline the failure paths deviate from).
+# --------------------------------------------------------------------------- #
+
+
+def test_clean_run_matches_stateless_reference_and_is_billable():
+    "A clean stream reproduces the stateless tokens and records completed work.\n\nTags: node, failure, billing"
+    model = _KvDenseLlama()
+    adapter = _make_adapter(model)
+    runner = HardenedSessionRunner(adapter)
+    prompt = [1, 2, 3, 4]
+    n_new = 8
+    outcome = runner.run(_generation("clean", prompt, n_new))
+    assert outcome.status is WorkStatus.COMPLETED
+    assert list(outcome.tokens) == model.stateless_greedy(prompt, n_new)
+    record = runner.work_ledger.records_for("clean")[0]
+    assert record.billable
+    assert record.tokens == n_new
+    assert runner.work_ledger.billable_tokens() == n_new
+
+
+# --------------------------------------------------------------------------- #
+# Deadlines and heartbeat/health loss terminate blocked operations.
+# --------------------------------------------------------------------------- #
+
+
+def test_deadline_terminates_a_blocked_stream_and_releases_kv():
+    "A deadline reached mid-stream terminates the run and frees its KV.\n\nTags: node, failure, deadline"
+    clock = _FakeClock()
+    adapter = _make_adapter()
+    manager = adapter.manager
+    runner = HardenedSessionRunner(adapter, clock=clock)
+
+    # Each step advances the clock by 1.0; the deadline fires at t=3.
+    def before_step(_step):
+        clock.advance(1.0)
+
+    outcome = runner.run(
+        _generation("slow", [5, 6, 7], 20),
+        deadline=3.0,
+        before_step=before_step,
+    )
+    assert outcome.status is WorkStatus.FAILED
+    assert outcome.failure_kind is FailureKind.DEADLINE_EXCEEDED
+    # The stream did not hang and did not finish: only the steps before the
+    # deadline committed, and the session's KV was released.
+    assert outcome.token_count < 20
+    assert isinstance(manager.resolve("slow", 0), CacheMiss)
+
+
+def test_heartbeat_loss_terminates_a_blocked_stream():
+    "Losing the peer heartbeat past the timeout terminates the stream.\n\nTags: node, failure, heartbeat"
+    clock = _FakeClock()
+    adapter = _make_adapter()
+    runner = HardenedSessionRunner(adapter, clock=clock)
+
+    def before_step(_step):
+        clock.advance(1.0)
+
+    # Heartbeats stop arriving after step 2; with a timeout of 1.5 the gap grows
+    # past the bound and the stream is terminated (health loss).
+    def heartbeat(step):
+        return step < 2
+
+    outcome = runner.run(
+        _generation("hb", [9, 8, 7], 20),
+        heartbeat_timeout=1.5,
+        heartbeat=heartbeat,
+        before_step=before_step,
+    )
+    assert outcome.status is WorkStatus.FAILED
+    assert outcome.failure_kind is FailureKind.HEARTBEAT_LOST
+    assert outcome.token_count < 20
+
+
+def test_deadline_guard_reports_remaining_and_resets_on_heartbeat():
+    "The guard exposes remaining time and a heartbeat resets the health timer.\n\nTags: node, failure, deadline"
+    clock = _FakeClock()
+    guard = DeadlineGuard(deadline=10.0, heartbeat_timeout=2.0, clock=clock)
+    guard.start()
+    guard.check()
+    assert guard.remaining() == 10.0
+    clock.advance(1.5)
+    guard.heartbeat()  # health refreshed at t=1.5
+    clock.advance(1.0)  # gap since heartbeat is 1.0 < 2.0
+    guard.check()
+    clock.advance(2.5)  # gap since heartbeat is now 3.5 > 2.0
+    with pytest.raises(StreamTerminated) as exc:
+        guard.check()
+    assert exc.value.kind is FailureKind.HEARTBEAT_LOST
+
+
+# --------------------------------------------------------------------------- #
+# Cancellation propagates across shards and releases KV + queued buffers.
+# --------------------------------------------------------------------------- #
+
+
+def test_cancellation_token_terminates_stream_and_releases_kv():
+    "A client cancel mid-stream stops the run and releases the session KV.\n\nTags: node, failure, cancel"
+    adapter = _make_adapter()
+    manager = adapter.manager
+    token = CancellationToken()
+    runner = HardenedSessionRunner(adapter)
+
+    # Cancel after two steps have run.
+    def before_step(step):
+        if step == 2:
+            token.cancel("client-hangup")
+
+    outcome = runner.run(
+        _generation("cancelme", [1, 2, 3], 20),
+        cancel_token=token,
+        before_step=before_step,
+    )
+    assert outcome.status is WorkStatus.CANCELLED
+    assert outcome.failure_kind is FailureKind.CANCELLED
+    assert outcome.token_count == 2  # steps 0 and 1 committed before the cancel
+    assert isinstance(manager.resolve("cancelme", 0), CacheMiss)
+
+
+def test_shard_cancellation_group_releases_every_shard_and_queued_buffers():
+    "One cancel frees KV on every node-local shard and releases queued buffers.\n\nTags: node, failure, cancel"
+    model = _KvDenseLlama()
+    # Three node-local shards of the same route, each with its own KV manager.
+    managers = []
+    for start, end in ((0, 1), (2, 3), (4, 5)):
+        shard = _KvReferenceShard(model, start, end)
+        mgr = HotKvStateManager(kv_recipe_for(shard))
+        mgr.open("route", 0)  # each holds live state for the session
+        managers.append(mgr)
+
+    released_buffers = []
+    group = ShardCancellationGroup("route", 0)
+    for mgr in managers:
+        group.add_shard(mgr)
+    group.add_queued_buffer(lambda: released_buffers.append("bundle-a"))
+    group.add_queued_buffer(lambda: released_buffers.append("bundle-b"))
+
+    outcome = group.cancel()
+    assert outcome.shards_released == 3
+    assert outcome.buffers_released == 2
+    assert released_buffers == ["bundle-a", "bundle-b"]
+    # Every shard's KV is gone: a lookup now yields an explicit released miss.
+    for mgr in managers:
+        miss = mgr.resolve("route", 0)
+        assert isinstance(miss, CacheMiss)
+        assert miss.reason is CacheMissReason.RELEASED
+    # Cancellation is idempotent.
+    again = group.cancel()
+    assert again.shards_released == 0
+    assert again.buffers_released == 0
+
+
+def test_scheduler_cancel_drains_queue_and_releases_active_kv():
+    "The scheduler cancel drops queued work and frees an active session's KV.\n\nTags: node, scheduler, cancel"
+    model = _KvDenseLlama()
+    shard = _KvReferenceShard(model, 0, model.n_layers - 1)
+    manager = HotKvStateManager(kv_recipe_for(shard))
+    engine = KvBatchEngine(KvBoundaryAdapter(shard, manager))
+    scheduler = ContinuousBatchScheduler(
+        engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4)
+    )
+    assert scheduler.submit(_generation("active", [1, 2, 3], 8)).running
+    assert scheduler.submit(_generation("waiting", [4, 5, 6], 8)).reason.value == "queued"
+    scheduler.run_tick()  # 'active' prefills and starts decoding, holding KV
+
+    # Cancel the queued one: it leaves the queue without ever taking a slot.
+    assert scheduler.cancel("waiting") is True
+    # Cancel the active one: its KV is released and it is recorded as cancelled.
+    assert scheduler.cancel("active") is True
+    assert manager.total_bytes == 0
+
+    telem = scheduler.telemetry()
+    assert telem.cancelled_sessions == 2
+    assert telem.completed_sessions == 0
+    assert telem.active_sessions == 0
+    assert telem.queue_depth == 0
+    # Cancelling an unknown / already-finished session is a no-op.
+    assert scheduler.cancel("active") is False
+    assert scheduler.cancel("never-seen") is False
+
+
+def test_scheduler_cancel_rejects_a_completed_reason():
+    "cancel() refuses a non-terminal reason so completed work is never faked.\n\nTags: node, scheduler, cancel"
+    model = _KvDenseLlama()
+    shard = _KvReferenceShard(model, 0, model.n_layers - 1)
+    manager = HotKvStateManager(kv_recipe_for(shard))
+    engine = KvBatchEngine(KvBoundaryAdapter(shard, manager))
+    scheduler = ContinuousBatchScheduler(engine)
+    scheduler.submit(_generation("x", [1, 2], 4))
+    with pytest.raises(Exception):
+        scheduler.cancel("x", reason=DoneReason.COMPLETED)
+
+
+# --------------------------------------------------------------------------- #
+# Idempotency: duplicate steps are no-ops; uncertain mutations never replay.
+# --------------------------------------------------------------------------- #
+
+
+def test_duplicate_step_delivery_is_idempotent_no_remutation():
+    "Replaying a committed step returns the recorded token without re-mutating KV.\n\nTags: node, failure, idempotency"
+    ledger = IdempotencyLedger()
+    key = StepKey("s", 0, 5)
+    disposition = ledger.begin(key)
+    assert disposition.fresh
+    ledger.commit(key, 42)
+    # A duplicate delivery of the same step returns the recorded token and is a
+    # no-op — the caller must not re-run the mutation.
+    replay = ledger.begin(key)
+    assert replay.duplicate
+    assert replay.token == 42
+
+
+def test_idempotent_run_replays_tokens_without_advancing_kv():
+    "Re-running a completed stream on the same ledger/epoch re-mutates nothing.\n\nTags: node, failure, idempotency"
+    model = _KvDenseLlama()
+    adapter = _make_adapter(model)
+    ledger = IdempotencyLedger()
+    runner = HardenedSessionRunner(adapter, idempotency=ledger)
+    request = _generation("idem", [3, 1, 4], 6)
+
+    first = runner.run(request)
+    assert first.status is WorkStatus.COMPLETED
+    kv_len_after_first = adapter.manager.get("idem", 0).seq_len
+
+    # A duplicate delivery of the entire stream: every step is a committed
+    # duplicate, so the runner replays the identical tokens and the KV length is
+    # unchanged (no double-append).
+    second = runner.run(request)
+    assert second.status is WorkStatus.COMPLETED
+    assert list(second.tokens) == list(first.tokens)
+    assert adapter.manager.get("idem", 0).seq_len == kv_len_after_first
+
+
+def test_uncertain_mutation_is_never_replayed_silently():
+    "A step marked uncertain refuses a silent replay; it must be verified/restarted.\n\nTags: node, failure, idempotency"
+    ledger = IdempotencyLedger()
+    key = StepKey("s", 0, 3)
+    ledger.begin(key)
+    ledger.mark_uncertain(key, "worker died before ack")
+    # Replaying an uncertain mutation is refused rather than silently re-applied.
+    with pytest.raises(UncertainMutationError):
+        ledger.begin(key)
+    assert ledger.has_uncertain()
+
+
+def test_in_flight_duplicate_is_treated_as_uncertain():
+    "A second begin before commit is refused (concurrent duplicate is unverified).\n\nTags: node, failure, idempotency"
+    ledger = IdempotencyLedger()
+    key = StepKey("s", 0, 1)
+    ledger.begin(key)  # in-flight, not yet committed
+    with pytest.raises(UncertainMutationError):
+        ledger.begin(key)
+
+
+# --------------------------------------------------------------------------- #
+# Worker death, stream reset, malformed bundle, stale epoch, cache miss.
+# --------------------------------------------------------------------------- #
+
+
+def test_worker_death_midstream_is_unverified_and_marks_step_uncertain():
+    "A worker dying mid-step yields unverified work and an unreplayable step.\n\nTags: node, failure, worker-death"
+    model = _KvDenseLlama()
+    # Fail on the 3rd step call (step index 2), after two tokens committed.
+    shard = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3)
+    adapter = _make_adapter(model, shard=shard)
+    ledger = IdempotencyLedger()
+    runner = HardenedSessionRunner(adapter, idempotency=ledger)
+
+    outcome = runner.run(_generation("dead", [1, 2, 3], 8))
+    assert outcome.status is WorkStatus.UNVERIFIED
+    assert outcome.failure_kind is FailureKind.WORKER_DEATH
+    assert outcome.token_count == 2  # the two committed steps
+    assert not outcome.completed
+    # The failed step is uncertain and can never be silently replayed.
+    assert ledger.has_uncertain()
+    with pytest.raises(UncertainMutationError):
+        ledger.begin(StepKey("dead", 0, 2))
+    # KV was released on failure.
+    assert isinstance(adapter.manager.resolve("dead", 0), CacheMiss)
+
+
+def test_stream_reset_is_restartable_failure():
+    "A stream reset injected mid-stream fails the run as a restartable transport loss.\n\nTags: node, failure, stream-reset"
+    adapter = _make_adapter()
+    runner = HardenedSessionRunner(adapter)
+
+    def before_step(step):
+        if step == 2:
+            raise StreamTerminated(FailureKind.STREAM_RESET, "peer reset the stream")
+
+    outcome = runner.run(_generation("reset", [1, 2, 3], 8), before_step=before_step)
+    assert outcome.status is WorkStatus.FAILED
+    assert outcome.failure_kind is FailureKind.STREAM_RESET
+    assert outcome.restartable
+
+
+def test_malformed_bundle_is_classified_and_does_not_corrupt_kv():
+    "A malformed activation bundle is rejected and leaves the KV context empty.\n\nTags: node, failure, malformed-bundle"
+    model = _KvDenseLlama()
+    mid = _KvReferenceShard(model, 2, 3)  # middle range: not head, not tail
+    manager = HotKvStateManager(kv_recipe_for(mid))
+    adapter = KvBoundaryAdapter(mid, manager)
+    assert not adapter.is_head and not adapter.is_tail
+
+    # A bundle that hands over at the wrong layer is malformed.
+    bad = BoundaryBundle(
+        architecture_adapter=adapter.architecture.adapter,
+        schema_version=adapter.architecture.boundary_schema_version,
+        tensor_name=adapter.architecture.boundary_tensor_name,
+        residual=np.zeros((1, 3, model.hidden), dtype=np.float32),
+        positions=np.arange(3, dtype=np.int64)[None, :],
+        next_layer=adapter.start_layer + 5,  # wrong handover layer
+        normalized=False,
+    )
+    with pytest.raises(BoundaryContractError) as exc:
+        adapter.prefill("mal", 0, boundary=bad)
+    assert classify_exception(exc.value) is FailureKind.MALFORMED_BUNDLE
+    # The malformed step never appended KV: the context is empty, not corrupted.
+    assert manager.get("mal", 0).seq_len == 0
+
+
+def test_stale_epoch_reference_is_rejected_and_classified():
+    "A reference to a superseded epoch is rejected as stale, never silently reused.\n\nTags: node, failure, stale-epoch"
+    model = _KvDenseLlama()
+    adapter = _make_adapter(model)
+    manager = adapter.manager
+    manager.open("sess", 5)  # current epoch is now 5
+    with pytest.raises(StaleRouteEpochError) as exc:
+        manager.resolve("sess", 4)  # epoch 4 is stale
+    assert classify_exception(exc.value) is FailureKind.STALE_EPOCH
+
+    # Driving the hardened runner on the stale epoch fails closed as STALE_EPOCH.
+    runner = HardenedSessionRunner(adapter)
+    outcome = runner.run(_generation("sess", [1, 2, 3], 4, epoch=3))
+    assert outcome.status is WorkStatus.FAILED
+    assert outcome.failure_kind is FailureKind.STALE_EPOCH
+
+
+def test_cache_miss_midstream_is_restartable():
+    "A KV eviction mid-stream surfaces an explicit cache miss the head can restart.\n\nTags: node, failure, cache-miss"
+    adapter = _make_adapter()
+    manager = adapter.manager
+    runner = HardenedSessionRunner(adapter)
+
+    # Evict the session's KV just before step 3's decode.
+    def before_step(step):
+        if step == 3:
+            manager.release("evict", 0)
+
+    outcome = runner.run(_generation("evict", [1, 2, 3], 10), before_step=before_step)
+    assert outcome.failure_kind is FailureKind.CACHE_MISS
+    assert outcome.restartable
+    assert outcome.token_count == 3  # steps 0..2 committed before the eviction
+
+
+# --------------------------------------------------------------------------- #
+# Alpha failover: restart from token zero, never import unverified KV.
+# --------------------------------------------------------------------------- #
+
+
+def test_alpha_failover_restarts_from_token_zero_and_completes():
+    "A transient worker death fails over to a fresh epoch and reproduces the tokens.\n\nTags: node, failure, failover"
+    model = _KvDenseLlama()
+    # Die on the 3rd step of the first attempt; the persistent call counter means
+    # the restart (which keeps counting) does not re-trip the fault.
+    shard = _FaultyShard(model, 0, model.n_layers - 1, fail_at_call=3)
+    adapter = _make_adapter(model, shard=shard)
+    manager = adapter.manager
+    runner = HardenedSessionRunner(adapter)
+    controller = RestartController([manager])
+
+    prompt = [7, 3, 9, 1]
+    n_new = 6
+    result = runner.run_with_failover(
+        _generation("alpha", prompt, n_new, epoch=0), controller, max_restarts=2
+    )
+    assert result.completed
+    assert result.restarts == 1
+    # The restart began on a fresh epoch and reproduced the full stateless stream
+    # from token zero — no half-computed KV was imported.
+    assert result.outcome.route_epoch == 1
+    assert list(result.outcome.tokens) == model.stateless_greedy(prompt, n_new)
+    # The failed epoch's KV is gone and the epoch is now stale.
+    with pytest.raises(StaleRouteEpochError):
+        manager.resolve("alpha", 0)
+    # First attempt was unverified, the restart completed: only the restart bills.
+    statuses = [a.status for a in result.attempts]
+    assert statuses == [WorkStatus.UNVERIFIED, WorkStatus.COMPLETED]
+    assert runner.work_ledger.billable_tokens() == n_new
+
+
+def test_failover_refuses_to_import_unverified_kv():
+    "assert_fresh_start fails closed if any shard still holds new-epoch KV.\n\nTags: node, failure, failover"
+    model = _KvDenseLlama()
+    adapter = _make_adapter(model)
+    manager = adapter.manager
+    controller = RestartController([manager])
+
+    new_epoch = controller.failover("s", 0)
+    assert new_epoch == 1
+    # A clean fresh start passes.
+    controller.assert_fresh_start("s", new_epoch)
+    # If unverified KV were present under the new epoch, the guard refuses it.
+    manager.open("s", new_epoch)
+    manager.append(
+        "s",
+        new_epoch,
+        {i: (np.zeros((1, model.n_heads, model.head_dim), dtype=np.float32),
+             np.zeros((1, model.n_heads, model.head_dim), dtype=np.float32))
+         for i in range(model.n_layers)},
+    )
+    with pytest.raises(Exception):
+        controller.assert_fresh_start("s", new_epoch)
+
+
+def test_non_restartable_failure_is_not_retried():
+    "A deterministic failure (deadline) returns immediately without a restart.\n\nTags: node, failure, failover"
+    clock = _FakeClock()
+    adapter = _make_adapter()
+    runner = HardenedSessionRunner(adapter, clock=clock)
+    controller = RestartController([adapter.manager])
+
+    def before_step(_step):
+        clock.advance(1.0)
+
+    result = runner.run_with_failover(
+        _generation("bounded", [1, 2, 3], 20),
+        controller,
+        max_restarts=3,
+        deadline=2.0,
+        before_step=before_step,
+    )
+    assert not result.completed
+    assert result.restarts == 0
+    assert result.outcome.failure_kind is FailureKind.DEADLINE_EXCEEDED
+
+
+# --------------------------------------------------------------------------- #
+# Billing / work records distinguish completed, cancelled, failed, unverified.
+# --------------------------------------------------------------------------- #
+
+
+def test_work_ledger_distinguishes_all_four_statuses():
+    "The work ledger keeps completed/cancelled/failed/unverified distinct.\n\nTags: node, failure, billing"
+    ledger = WorkLedger()
+    ledger.record(WorkRecord("a", 0, WorkStatus.COMPLETED, tokens=8))
+    ledger.record(WorkRecord("b", 0, WorkStatus.CANCELLED, tokens=3,
+                             failure_kind=FailureKind.CANCELLED))
+    ledger.record(WorkRecord("c", 0, WorkStatus.FAILED, tokens=1,
+                             failure_kind=FailureKind.DEADLINE_EXCEEDED))
+    ledger.record(WorkRecord("d", 0, WorkStatus.UNVERIFIED, tokens=2,
+                             failure_kind=FailureKind.WORKER_DEATH))
+
+    counts = ledger.counts_by_status()
+    assert counts == {
+        "completed": 1, "cancelled": 1, "failed": 1, "unverified": 1,
+    }
+    # Only completed work is billable — cancelled/failed/unverified tokens are
+    # recorded for observability but never charged.
+    assert ledger.billable_tokens() == 8
+    assert [r.session_id for r in ledger.billable_records()] == ["a"]
+    # JSON-safe for durable evidence.
+    payload = ledger.to_dict()
+    assert payload["billable_tokens"] == 8
+    assert payload["counts_by_status"]["unverified"] == 1
+    json.dumps(payload)
+
+
+def test_work_status_and_classification_mapping():
+    "Failure kinds map to the right billing status and exception classes.\n\nTags: node, failure, billing"
+    assert work_status_for(FailureKind.CANCELLED) is WorkStatus.CANCELLED
+    assert work_status_for(FailureKind.WORKER_DEATH) is WorkStatus.UNVERIFIED
+    # A stream reset detected at a step boundary is a certain failure (nothing
+    # committed for that step) — only an unexpected mid-step error is unverified.
+    assert work_status_for(FailureKind.STREAM_RESET) is WorkStatus.FAILED
+    assert work_status_for(FailureKind.DEADLINE_EXCEEDED) is WorkStatus.FAILED
+    assert work_status_for(FailureKind.MALFORMED_BUNDLE) is WorkStatus.FAILED
+    assert work_status_for(FailureKind.STALE_EPOCH) is WorkStatus.FAILED
+    assert work_status_for(FailureKind.CACHE_MISS) is WorkStatus.FAILED
+
+    assert classify_exception(OperationCancelled()) is FailureKind.CANCELLED
+    assert classify_exception(StaleRouteEpochError("x")) is FailureKind.STALE_EPOCH
+    assert classify_exception(BoundaryContractError("x")) is FailureKind.MALFORMED_BUNDLE
+    assert classify_exception(RuntimeError("boom")) is FailureKind.WORKER_DEATH
+    assert (
+        classify_exception(StreamTerminated(FailureKind.HEARTBEAT_LOST))
+        is FailureKind.HEARTBEAT_LOST
+    )