feat: checkpoint batching and release-gate stories

tests/test_batch_scheduler.py

@@ -0,0 +1,472 @@
+"""Continuous batching and bounded admission (DGR-012).
+
+These tests drive the node-local continuous-batching scheduler with the *same*
+pure-numpy KV-cached dense-Llama reference the Hot KV State manager uses
+(DGR-007), imported from ``test_hot_kv_state``. That keeps the whole gate
+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:
+
+* bounded admission against weight/KV/scratch/queue budgets,
+* compatible decode steps batched with per-session positions/outputs preserved,
+* prefill never starving in-flight decode (explicit decode-first policy),
+* backpressure when the bounded queue is full,
+* capability telemetry reporting every required signal,
+* a deterministic 1/2/4/8 concurrency sweep showing saturation and no
+  cross-session corruption.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from meshnet_node.hot_kv_state import (
+    HotKvStateConfig,
+    HotKvStateManager,
+    KvBoundaryAdapter,
+    kv_recipe_for,
+)
+from meshnet_node.batch_scheduler import (
+    AdmissionReason,
+    ContinuousBatchScheduler,
+    GenerationRequest,
+    KvBatchEngine,
+    NodeBudget,
+    Phase,
+    run_concurrency_sweep,
+)
+
+# 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
+
+
+def _make_engine(
+    model: _KvDenseLlama | None = None,
+    *,
+    config: HotKvStateConfig | None = None,
+) -> KvBatchEngine:
+    """A full-shard KV batch engine over the deterministic numpy dense-Llama."""
+    model = model or _KvDenseLlama()
+    shard = _KvReferenceShard(model, 0, model.n_layers - 1)
+    manager = HotKvStateManager(kv_recipe_for(shard), config=config)
+    adapter = KvBoundaryAdapter(shard, manager)
+    return KvBatchEngine(adapter)
+
+
+def _reference_tokens(model: _KvDenseLlama, prompt, n_new: int) -> list[int]:
+    return model.stateless_greedy(list(prompt), n_new)
+
+
+def _generation(session_id: str, prompt, n_new: int, epoch: int = 0) -> GenerationRequest:
+    return GenerationRequest(
+        session_id=session_id,
+        route_epoch=epoch,
+        prompt_token_ids=tuple(prompt),
+        max_new_tokens=n_new,
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Bounded admission (weight / KV / scratch / queue budgets).
+# --------------------------------------------------------------------------- #
+
+
+def test_admission_respects_active_scratch_and_queue_budgets():
+    "Admission fills active slots, queues the overflow, then rejects a full queue.\n\nTags: node, scheduler, admission"
+    engine = _make_engine()
+    budget = NodeBudget(
+        max_active_sessions=2,
+        scratch_bytes_per_session=1,
+        scratch_budget_bytes=2,  # scratch also caps at 2 concurrent
+        max_queue_depth=1,
+        max_batch_size=2,
+    )
+    scheduler = ContinuousBatchScheduler(engine, budget)
+
+    a = scheduler.submit(_generation("a", [1, 2, 3], 4))
+    b = scheduler.submit(_generation("b", [4, 5, 6], 4))
+    assert a.reason is AdmissionReason.ADMITTED
+    assert b.reason is AdmissionReason.ADMITTED
+
+    # Two active slots full -> the next goes to the bounded queue.
+    c = scheduler.submit(_generation("c", [7, 8, 9], 4))
+    assert c.reason is AdmissionReason.QUEUED
+
+    # Queue depth 1 is now full -> backpressure rejection.
+    d = scheduler.submit(_generation("d", [1, 1, 1], 4))
+    assert d.reason is AdmissionReason.REJECTED_QUEUE_FULL
+    assert d.rejected
+
+    telem = scheduler.telemetry()
+    assert telem.active_sessions == 2
+    assert telem.queue_depth == 1
+    assert telem.rejected_admissions_total == 1
+    assert telem.rejected_by_reason[AdmissionReason.REJECTED_QUEUE_FULL.value] == 1
+
+
+def test_admission_rejects_a_session_that_cannot_fit_the_kv_budget():
+    "A generation whose whole KV cannot fit the node budget is rejected up front.\n\nTags: node, scheduler, admission"
+    engine = _make_engine()
+    per_token = engine._manager.recipe.bytes_per_token()
+    # Budget holds only 3 positions; a prompt(4)+7 new = 10 final positions cannot fit.
+    budget = NodeBudget(kv_budget_bytes=per_token * 3)
+    scheduler = ContinuousBatchScheduler(engine, budget)
+    decision = scheduler.submit(_generation("big", [1, 2, 3, 4], 7))
+    assert decision.reason is AdmissionReason.REJECTED_KV_BUDGET
+    assert scheduler.telemetry().rejected_admissions_total == 1
+
+
+def test_admission_rejects_when_per_session_scratch_exceeds_budget():
+    "A per-session scratch larger than the whole scratch envelope is rejected.\n\nTags: node, scheduler, admission"
+    engine = _make_engine()
+    budget = NodeBudget(scratch_bytes_per_session=1024, scratch_budget_bytes=512)
+    scheduler = ContinuousBatchScheduler(engine, budget)
+    decision = scheduler.submit(_generation("s", [1, 2], 2))
+    assert decision.reason is AdmissionReason.REJECTED_SCRATCH_BUDGET
+
+
+def test_duplicate_submission_is_rejected():
+    "Submitting a session id that is already scheduled is rejected as a duplicate.\n\nTags: node, scheduler, admission"
+    engine = _make_engine()
+    scheduler = ContinuousBatchScheduler(engine, NodeBudget(max_active_sessions=4))
+    assert scheduler.submit(_generation("dup", [1, 2], 3)).reason is AdmissionReason.ADMITTED
+    assert scheduler.submit(_generation("dup", [3, 4], 3)).reason is AdmissionReason.REJECTED_DUPLICATE
+
+
+def test_weight_budget_is_reported_in_telemetry():
+    "The resident weight footprint is surfaced as a capability signal.\n\nTags: node, scheduler, telemetry"
+    engine = _make_engine()
+    budget = NodeBudget(weight_bytes=123_456)
+    scheduler = ContinuousBatchScheduler(engine, budget)
+    assert scheduler.telemetry().weight_bytes == 123_456
+
+
+# --------------------------------------------------------------------------- #
+# Continuous batching preserves per-session positions and outputs.
+# --------------------------------------------------------------------------- #
+
+
+def test_batched_decode_preserves_per_session_positions_and_outputs():
+    "Four sessions batched together each reproduce their own stateless tokens.\n\nTags: node, scheduler, batching"
+    model = _KvDenseLlama()
+    engine = _make_engine(model)
+    budget = NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4)
+    scheduler = ContinuousBatchScheduler(engine, budget)
+
+    prompts = {
+        "alpha": [1, 2, 3, 4],
+        "bravo": [40, 39, 2, 15],
+        "charlie": [7, 7, 7, 7],
+        "delta": [31, 5, 18, 22],
+    }
+    n_new = 10
+    references = {sid: _reference_tokens(model, p, n_new) for sid, p in prompts.items()}
+    # The four references must diverge, else "no cross-talk" would be vacuous.
+    assert len({tuple(v) for v in references.values()}) == 4
+
+    for sid, prompt in prompts.items():
+        assert scheduler.submit(_generation(sid, prompt, n_new)).running
+
+    outputs = scheduler.run_to_completion()
+    for sid in prompts:
+        assert outputs[sid] == references[sid], sid
+
+    telem = scheduler.telemetry()
+    # A genuine batch formed: at least one decode tick carried all four sessions.
+    assert telem.batch_occupancy_max == 4
+    assert telem.completed_sessions == 4
+    assert telem.active_sessions == 0
+
+
+def test_positions_are_isolated_across_different_prompt_lengths():
+    "Sessions with different prompt lengths keep independent positions when batched.\n\nTags: node, scheduler, batching"
+    model = _KvDenseLlama()
+    engine = _make_engine(model)
+    scheduler = ContinuousBatchScheduler(
+        engine, NodeBudget(max_active_sessions=3, max_batch_size=3, max_queue_depth=3)
+    )
+    jobs = {
+        "short": ([5], 6),
+        "medium": ([2, 9, 14], 6),
+        "long": ([1, 2, 3, 4, 5, 6, 7], 6),
+    }
+    refs = {sid: _reference_tokens(model, p, n) for sid, (p, n) in jobs.items()}
+    for sid, (prompt, n) in jobs.items():
+        scheduler.submit(_generation(sid, prompt, n))
+    outputs = scheduler.run_to_completion()
+    for sid in jobs:
+        assert outputs[sid] == refs[sid], sid
+
+
+# --------------------------------------------------------------------------- #
+# Prefill does not starve decode.
+# --------------------------------------------------------------------------- #
+
+
+def test_prefill_does_not_starve_in_flight_decode():
+    "A burst of new prefills never stalls an already-decoding session.\n\nTags: node, scheduler, fairness"
+    model = _KvDenseLlama()
+    engine = _make_engine(model)
+    # One prefill per tick (budget == a single prompt) so prefill is throttled and
+    # we can observe that decode still advances every tick.
+    budget = NodeBudget(
+        max_active_sessions=8,
+        max_batch_size=8,
+        max_queue_depth=8,
+        scratch_bytes_per_session=1,
+        scratch_budget_bytes=8,
+        max_prefill_tokens_per_tick=4,
+    )
+    scheduler = ContinuousBatchScheduler(engine, budget)
+
+    # Session A starts and prefills on tick 1.
+    scheduler.submit(_generation("A", [3, 14, 1, 5], 12))
+    scheduler.run_tick()
+    a_state = scheduler.session_result("A")
+    assert a_state.phase is Phase.DECODING
+    a_len = len(a_state.generated)
+    assert a_len == 1
+
+    # Burst of new work arrives while A is decoding.
+    for sid in ("B", "C", "D", "E"):
+        scheduler.submit(_generation(sid, [2, 27, 18, 4], 12))
+
+    # Over the next few ticks A must decode on *every* tick (never starved),
+    # while at most one new session prefills per tick (prefill is bounded).
+    prefill_counts = []
+    for _ in range(4):
+        report = scheduler.run_tick()
+        new_a_len = len(scheduler.session_result("A").generated)
+        assert new_a_len == a_len + 1, "decode of A stalled while prefills were pending"
+        a_len = new_a_len
+        assert "A" in report.decoded
+        prefill_counts.append(len(report.prefilled))
+
+    assert max(prefill_counts) <= 1, "prefill was not bounded per tick"
+
+
+def test_decode_first_policy_is_explicit_in_a_single_tick():
+    "In one tick decode of active sessions precedes prefill of new ones.\n\nTags: node, scheduler, fairness"
+    model = _KvDenseLlama()
+    engine = _make_engine(model)
+    scheduler = ContinuousBatchScheduler(
+        engine,
+        NodeBudget(max_active_sessions=4, max_batch_size=4, max_queue_depth=4,
+                   scratch_bytes_per_session=1, scratch_budget_bytes=4),
+    )
+    scheduler.submit(_generation("live", [1, 2, 3], 8))
+    scheduler.run_tick()  # 'live' prefills, now decoding
+    scheduler.submit(_generation("fresh", [9, 8, 7], 8))
+    report = scheduler.run_tick()
+    assert "live" in report.decoded
+    assert "fresh" in report.prefilled
+
+
+# --------------------------------------------------------------------------- #
+# Backpressure and bounded memory.
+# --------------------------------------------------------------------------- #
+
+
+def test_backpressure_signals_when_queue_full_then_recovers():
+    "A full queue rejects new work; a completed session frees a slot for the queue.\n\nTags: node, scheduler, backpressure"
+    engine = _make_engine()
+    budget = NodeBudget(
+        max_active_sessions=1,
+        max_batch_size=1,
+        max_queue_depth=1,
+        scratch_bytes_per_session=1,
+        scratch_budget_bytes=1,
+    )
+    scheduler = ContinuousBatchScheduler(engine, budget)
+    assert scheduler.submit(_generation("first", [1, 2], 2)).running
+    assert scheduler.submit(_generation("second", [3, 4], 2)).reason is AdmissionReason.QUEUED
+    # Both a slot and the queue are full now.
+    assert scheduler.submit(_generation("third", [5, 6], 2)).reason is AdmissionReason.REJECTED_QUEUE_FULL
+
+    # Drain 'first'; the queued 'second' must be pulled into the freed slot.
+    scheduler.run_to_completion()
+    outputs = scheduler.outputs()
+    assert set(outputs) == {"first", "second"}
+
+
+def test_completed_sessions_release_kv_so_growth_is_bounded():
+    "Finished sessions release their KV, so total KV returns to zero.\n\nTags: node, scheduler, backpressure"
+    engine = _make_engine()
+    scheduler = ContinuousBatchScheduler(
+        engine, NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=8)
+    )
+    for sid in ("a", "b", "c", "d"):
+        scheduler.submit(_generation(sid, [1, 2, 3], 4))
+    scheduler.run_to_completion()
+    telem = scheduler.telemetry()
+    assert telem.kv_total_bytes == 0, "KV not released after completion"
+    assert telem.active_sessions == 0
+    assert telem.completed_sessions == 4
+
+
+# --------------------------------------------------------------------------- #
+# Telemetry.
+# --------------------------------------------------------------------------- #
+
+
+def test_telemetry_reports_every_required_signal():
+    "The capability snapshot reports sessions, queue, batch, KV, rates, rejections.\n\nTags: node, scheduler, telemetry"
+    model = _KvDenseLlama()
+    engine = _make_engine(model)
+    clock = _FakeClock()
+    budget = NodeBudget(max_active_sessions=2, max_batch_size=2, max_queue_depth=1)
+    scheduler = ContinuousBatchScheduler(engine, budget, clock=clock)
+
+    scheduler.submit(_generation("x", [1, 2, 3], 4))
+    scheduler.submit(_generation("y", [4, 5, 6], 4))
+    scheduler.submit(_generation("z", [7, 8, 9], 4))  # queued
+    rejected = scheduler.submit(_generation("w", [1, 1, 1], 4))  # queue full
+    assert rejected.rejected
+
+    clock.advance(1.0)
+    scheduler.run_tick()  # both prefill
+    clock.advance(1.0)
+    scheduler.run_tick()  # both decode as a batch of 2
+
+    clock.advance(2.0)
+    telem = scheduler.telemetry()
+    snap = telem.to_dict()
+    for key in (
+        "active_sessions", "queue_depth", "batch_occupancy_last",
+        "batch_occupancy_avg", "batch_occupancy_max", "weight_bytes",
+        "kv_total_bytes", "kv_budget_bytes", "kv_pressure",
+        "scratch_used_bytes", "scratch_budget_bytes", "scratch_pressure",
+        "prefill_tokens_total", "decode_tokens_total",
+        "prefill_tokens_per_sec", "decode_tokens_per_sec",
+        "rejected_admissions_total", "rejected_by_reason",
+        "completed_sessions", "ticks",
+    ):
+        assert key in snap, key
+
+    assert telem.batch_occupancy_max == 2
+    assert telem.prefill_tokens_total == 6  # two prompts of length 3
+    assert telem.decode_tokens_total == 2  # one batched decode step, two sessions
+    assert telem.rejected_admissions_total == 1
+    # Rates are deterministic under the injected clock: 4 seconds elapsed.
+    assert telem.decode_tokens_per_sec == pytest.approx(2 / 4.0)
+    assert telem.prefill_tokens_per_sec == pytest.approx(6 / 4.0)
+    assert 0.0 < telem.kv_pressure <= 1.0
+
+
+# --------------------------------------------------------------------------- #
+# Concurrency 1/2/4/8 sweep: saturation and no corruption.
+# --------------------------------------------------------------------------- #
+
+
+def test_concurrency_sweep_identifies_saturation_without_corruption():
+    "A 1/2/4/8 sweep raises batch occupancy, cuts ticks, and never corrupts output.\n\nTags: node, scheduler, benchmark"
+    model = _KvDenseLlama()
+    prompts = {
+        "s0": [1, 2, 3, 4], "s1": [5, 6, 7, 8], "s2": [9, 10, 11, 12],
+        "s3": [13, 14, 15, 16], "s4": [17, 18, 19, 20], "s5": [21, 22, 23, 24],
+        "s6": [25, 26, 27, 28], "s7": [29, 30, 31, 32],
+    }
+    n_new = 8
+    requests = [_generation(sid, p, n_new) for sid, p in prompts.items()]
+
+    sweep = run_concurrency_sweep(
+        lambda: _make_engine(model),
+        requests,
+        concurrency_levels=(1, 2, 4, 8),
+    )
+
+    assert sweep.corruption_free
+    assert [r.concurrency for r in sweep.results] == [1, 2, 4, 8]
+
+    # No session hit a cache miss (budgets are sized to never evict here).
+    assert all(r.cache_misses == 0 for r in sweep.results)
+    assert all(r.rejected_admissions == 0 for r in sweep.results)
+
+    # Each per-session stream matches the serialized (concurrency-1) reference.
+    for sid, prompt in prompts.items():
+        assert list(sweep.reference_outputs[sid]) == _reference_tokens(model, prompt, n_new)
+
+    occupancies = [r.avg_batch_occupancy for r in sweep.results]
+    ticks = [r.ticks for r in sweep.results]
+    tokens_per_tick = [r.tokens_per_tick for r in sweep.results]
+
+    # Batching packs more sessions per decode step as concurrency rises, so
+    # average occupancy strictly increases and total ticks strictly decrease.
+    assert occupancies == sorted(occupancies) and len(set(occupancies)) == 4
+    assert ticks == sorted(ticks, reverse=True) and len(set(ticks)) == 4
+    # Aggregate work per tick rises with concurrency (the throughput win).
+    assert tokens_per_tick == sorted(tokens_per_tick)
+
+    # For eight equal-length jobs the node keeps saturating up to the top level.
+    assert sweep.saturation_concurrency == 8
+
+    # The report is JSON-safe for durable evidence.
+    import json
+
+    json.dumps(sweep.to_dict())
+
+
+def test_concurrency_sweep_saturates_below_max_when_load_is_small():
+    "With fewer concurrent jobs than slots, saturation is found below the top level.\n\nTags: node, scheduler, benchmark"
+    model = _KvDenseLlama()
+    # Only three jobs: at concurrency 4 and 8 the batch can never exceed 3, so
+    # occupancy stops rising past the load and saturation is detected early.
+    requests = [
+        _generation("j0", [1, 2, 3], 6),
+        _generation("j1", [4, 5, 6], 6),
+        _generation("j2", [7, 8, 9], 6),
+    ]
+    sweep = run_concurrency_sweep(
+        lambda: _make_engine(model), requests, concurrency_levels=(1, 2, 4, 8)
+    )
+    assert sweep.corruption_free
+    assert sweep.saturation_concurrency <= 4
+    # Levels at or above the load size share the same occupancy/tick profile.
+    top = [r for r in sweep.results if r.concurrency >= 4]
+    assert len({r.ticks for r in top}) == 1
+
+
+# --------------------------------------------------------------------------- #
+# Engine contract guards.
+# --------------------------------------------------------------------------- #
+
+
+def test_kv_batch_engine_requires_a_full_shard():
+    "The batch engine rejects a partial (non head+tail) shard.\n\nTags: node, scheduler"
+    model = _KvDenseLlama()
+    head = _KvReferenceShard(model, 0, 2)  # head only, not tail
+    manager = HotKvStateManager(kv_recipe_for(head))
+    adapter = KvBoundaryAdapter(head, manager)
+    with pytest.raises(Exception):
+        KvBatchEngine(adapter)
+
+
+def test_run_to_completion_is_bounded_against_misconfiguration():
+    "run_to_completion raises rather than looping forever when work cannot drain.\n\nTags: node, scheduler"
+    engine = _make_engine()
+    scheduler = ContinuousBatchScheduler(
+        engine, NodeBudget(max_active_sessions=1, max_batch_size=1, max_queue_depth=4)
+    )
+    scheduler.submit(_generation("only", [1, 2], 3))
+    # A tiny explicit tick ceiling is exceeded deterministically.
+    with pytest.raises(Exception):
+        scheduler.run_to_completion(max_ticks=1)

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
+    )