"""Bounded failure, cancellation, and restart semantics for Shard streams (DGR-013).

Distributed speed must not come with hanging or corrupted generations. This module
hardens the per-Route-Session decode stream that runs over the DGR-007 Hot KV State
manager (isolated ``(session, epoch)`` KV) and the DGR-012 continuous-batch
scheduler. It is deliberately backend-agnostic and pure-Python: it drives the same
``KvBoundaryAdapter`` the default deterministic gate uses, so the whole matrix stays
download-free, GPU-free, and API-credit-free while exercising the *real* KV
isolation path (the pinned llama.cpp worker, DGR-008, implements the identical
adapter contract).

The guarantees, mapped to the story's acceptance criteria:

* **Deadlines and heartbeat/health loss terminate blocked stream operations.**
  :class:`DeadlineGuard` bounds every step against an absolute deadline and a
  heartbeat-timeout; when either is breached it raises :class:`StreamTerminated`
  so a blocked stream never hangs.
* **Cancellation propagates across every Shard and releases local KV and queued
  buffers.** :class:`ShardCancellationGroup` fans a single cancel across every
  node-local KV manager serving a Route Session and releases queued activation
  buffers; the DGR-012 scheduler's :meth:`~meshnet_node.batch_scheduler.
  ContinuousBatchScheduler.cancel` drops queued/active work on this node.
* **Duplicate steps are idempotent; uncertain mutations are never replayed
  silently.** :class:`IdempotencyLedger` records each committed
  ``(session, epoch, step)`` and returns the recorded token for a duplicate
  delivery instead of re-running it. A step whose outcome is *uncertain* (the
  worker died mid-mutation) is marked uncertain and can never be silently
  replayed — a replay attempt raises :class:`UncertainMutationError`, forcing an
  explicit verify-or-restart.
* **Alpha failover restarts from token zero on a newly compatible route rather
  than importing unverified KV.** :class:`RestartController` opens a *new* route
  epoch, releases every shard's prior-epoch KV, and the restart re-prefills the
  whole prompt from token zero. The old epoch becomes stale (rejected by the KV
  manager); unverified KV is never migrated (RALPH runtime decision #14).
* **Billing/work records distinguish completed, cancelled, failed, and unverified
  work.** :class:`WorkLedger` records a typed :class:`WorkRecord` per attempt;
  only :attr:`WorkStatus.COMPLETED` records are billable, so cancelled, failed,
  and uncertain (unverified) work is accounted but never charged.

:class:`HardenedSessionRunner` composes these into one drivable stream: it runs a
single session's prefill+decode through the adapter under a deadline/heartbeat
guard and a cancellation token, records the typed work outcome, and — via
:meth:`HardenedSessionRunner.run_with_failover` — restarts a transient failure
from token zero on a fresh epoch.
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Any, Callable, Mapping, Sequence

from meshnet_node.batch_scheduler import DoneReason, GenerationRequest
from meshnet_node.boundary_adapter import BoundaryContractError, TailOutput
from meshnet_node.hot_kv_state import (
    CacheMiss,
    HotKvStateManager,
    IncompatibleCacheRecipeError,
    KvBoundaryAdapter,
    KvCacheMissError,
    StaleRouteEpochError,
)


class FailureSemanticsError(RuntimeError):
    """Base class for failure/cancellation/restart errors."""


# --------------------------------------------------------------------------- #
# Typed outcomes: failure kinds and billing/work statuses.
# --------------------------------------------------------------------------- #


class FailureKind(str, Enum):
    """Why a stream step failed. Stable strings for the protocol's structured status."""

    # Bounded termination of a blocked op.
    DEADLINE_EXCEEDED = "deadline-exceeded"
    HEARTBEAT_LOST = "heartbeat-lost"
    # Transport / worker loss (transient — a restart from token zero may succeed).
    WORKER_DEATH = "worker-death"
    STREAM_RESET = "stream-reset"
    # Protocol violations (deterministic — a restart would fail identically).
    MALFORMED_BUNDLE = "malformed-bundle"
    STALE_EPOCH = "stale-epoch"
    INCOMPATIBLE_RECIPE = "incompatible-recipe"
    # KV state expected by the caller is gone; re-prefill from token zero.
    CACHE_MISS = "cache-miss"
    # Explicit client cancellation.
    CANCELLED = "cancelled"


# Failure kinds that a from-token-zero restart on a fresh route may recover from.
# A protocol violation or an explicit bound (deadline/cancel) is NOT restartable —
# retrying it would hang or fail identically, so we surface it instead.
_RESTARTABLE = frozenset(
    {
        FailureKind.WORKER_DEATH,
        FailureKind.STREAM_RESET,
        FailureKind.CACHE_MISS,
    }
)

# Failure kinds whose mutation outcome is *uncertain* — the KV may or may not have
# advanced, so the confirmed work is billed as UNVERIFIED and never replayed
# silently. Only an *unexpected* error raised while a step was executing is
# uncertain (mapped to WORKER_DEATH). A stream reset, deadline, or cache miss
# detected at a step boundary is certain: nothing committed for that step.
_UNCERTAIN = frozenset({FailureKind.WORKER_DEATH})


class WorkStatus(str, Enum):
    """The billing-relevant outcome class of a unit of work (AC: billing records).

    Only :attr:`COMPLETED` work is billable. Cancelled, failed, and unverified
    work is recorded distinctly so a client is never charged for a generation that
    hung, was cancelled, or whose mutations could not be verified.
    """

    COMPLETED = "completed"
    CANCELLED = "cancelled"
    FAILED = "failed"
    UNVERIFIED = "unverified"


def work_status_for(kind: FailureKind) -> WorkStatus:
    """Map a terminal failure kind to its billing/work status."""
    if kind is FailureKind.CANCELLED:
        return WorkStatus.CANCELLED
    if kind in _UNCERTAIN:
        return WorkStatus.UNVERIFIED
    return WorkStatus.FAILED


def classify_exception(exc: BaseException) -> FailureKind:
    """Classify a raised error into a :class:`FailureKind`.

    Protocol violations map to their specific kind; a :class:`StreamTerminated`
    carries its own kind; any *unexpected* error is treated as worker death
    (an uncertain, transient loss), never silently ignored.
    """
    if isinstance(exc, StreamTerminated):
        return exc.kind
    if isinstance(exc, OperationCancelled):
        return FailureKind.CANCELLED
    if isinstance(exc, StaleRouteEpochError):
        return FailureKind.STALE_EPOCH
    if isinstance(exc, IncompatibleCacheRecipeError):
        return FailureKind.INCOMPATIBLE_RECIPE
    if isinstance(exc, BoundaryContractError):
        return FailureKind.MALFORMED_BUNDLE
    if isinstance(exc, KvCacheMissError):
        return FailureKind.CACHE_MISS
    return FailureKind.WORKER_DEATH


# --------------------------------------------------------------------------- #
# Deadlines and heartbeat/health loss.
# --------------------------------------------------------------------------- #


class StreamTerminated(FailureSemanticsError):
    """A blocked stream op was terminated by a deadline or heartbeat/health loss."""

    def __init__(self, kind: FailureKind, detail: str = "") -> None:
        self.kind = kind
        self.detail = detail
        suffix = f": {detail}" if detail else ""
        super().__init__(f"stream terminated ({kind.value}){suffix}")


class OperationCancelled(FailureSemanticsError):
    """Raised when a step observes its :class:`CancellationToken` is cancelled."""

    def __init__(self, reason: str = "client-cancel") -> None:
        self.reason = reason
        super().__init__(f"operation cancelled: {reason}")


@dataclass
class DeadlineGuard:
    """Bounds a blocked stream op against an absolute deadline and heartbeat loss.

    ``deadline`` is an absolute time on ``clock``'s scale (``None`` disables it).
    ``heartbeat_timeout`` is the maximum tolerated gap since the last observed
    heartbeat; when the peer stops sending heartbeats (its health is lost) the gap
    grows past the timeout and :meth:`check` raises rather than blocking forever.
    Both bounds are checked with an injected ``clock`` so the matrix is
    deterministic.
    """

    deadline: float | None = None
    heartbeat_timeout: float | None = None
    clock: Callable[[], float] = time.monotonic
    _last_heartbeat: float = field(default=0.0, init=False)
    _started: bool = field(default=False, init=False)

    def __post_init__(self) -> None:
        if self.heartbeat_timeout is not None and self.heartbeat_timeout <= 0:
            raise FailureSemanticsError("heartbeat_timeout must be positive")

    def start(self) -> None:
        self._last_heartbeat = self.clock()
        self._started = True

    def heartbeat(self) -> None:
        """Record that the peer is alive (resets the heartbeat gap)."""
        self._last_heartbeat = self.clock()

    def check(self) -> None:
        """Raise :class:`StreamTerminated` if the deadline or heartbeat lapsed."""
        if not self._started:
            self.start()
        now = self.clock()
        if self.deadline is not None and now >= self.deadline:
            raise StreamTerminated(
                FailureKind.DEADLINE_EXCEEDED,
                f"deadline {self.deadline} reached at {now}",
            )
        if self.heartbeat_timeout is not None:
            gap = now - self._last_heartbeat
            if gap > self.heartbeat_timeout:
                raise StreamTerminated(
                    FailureKind.HEARTBEAT_LOST,
                    f"no heartbeat for {gap} > {self.heartbeat_timeout}",
                )

    def remaining(self) -> float | None:
        if self.deadline is None:
            return None
        return self.deadline - self.clock()


# --------------------------------------------------------------------------- #
# Cancellation that propagates across shards and releases KV + queued buffers.
# --------------------------------------------------------------------------- #


class CancellationToken:
    """A thread-safe one-shot cancellation flag shared by a Route Session's steps."""

    def __init__(self) -> None:
        self._cancelled = False
        self._reason = ""
        self._lock = threading.Lock()

    def cancel(self, reason: str = "client-cancel") -> None:
        with self._lock:
            if not self._cancelled:
                self._cancelled = True
                self._reason = reason

    @property
    def cancelled(self) -> bool:
        with self._lock:
            return self._cancelled

    @property
    def reason(self) -> str:
        with self._lock:
            return self._reason

    def raise_if_cancelled(self) -> None:
        with self._lock:
            if self._cancelled:
                raise OperationCancelled(self._reason)


@dataclass(frozen=True)
class CancellationOutcome:
    """What a :meth:`ShardCancellationGroup.cancel` released (for observability)."""

    session_id: str
    route_epoch: int
    shards_released: int
    buffers_released: int

    def to_dict(self) -> dict:
        return {
            "session_id": self.session_id,
            "route_epoch": self.route_epoch,
            "shards_released": self.shards_released,
            "buffers_released": self.buffers_released,
        }


class ShardCancellationGroup:
    """Fan one cancellation across every node-local Shard of a Route Session.

    A Route Session spans a chain of Shards, each with its own local Hot KV State
    manager (KV is never migrated between nodes). Cancelling the session must free
    *all* of that state: this group releases the ``(session, epoch)`` KV on every
    registered manager and invokes every registered queued-buffer release callback
    (the pending activation bundles a node holds for the session). Release is
    idempotent, so cancelling twice is safe.
    """

    def __init__(self, session_id: str, route_epoch: int) -> None:
        if not isinstance(session_id, str) or not session_id.strip():
            raise FailureSemanticsError("session_id must be a non-empty string")
        self.session_id = session_id
        self.route_epoch = int(route_epoch)
        self._managers: list[HotKvStateManager] = []
        self._buffers: list[Callable[[], None]] = []
        self._lock = threading.Lock()
        self._cancelled = False

    def add_shard(self, manager: HotKvStateManager) -> "ShardCancellationGroup":
        with self._lock:
            self._managers.append(manager)
        return self

    def add_queued_buffer(
        self, release: Callable[[], None]
    ) -> "ShardCancellationGroup":
        """Register a queued activation buffer's release callback."""
        with self._lock:
            self._buffers.append(release)
        return self

    @property
    def cancelled(self) -> bool:
        with self._lock:
            return self._cancelled

    def cancel(self) -> CancellationOutcome:
        """Release every shard's KV and every queued buffer for this session."""
        with self._lock:
            managers = list(self._managers)
            buffers = list(self._buffers)
            self._buffers.clear()
            self._cancelled = True
        shards_released = 0
        for manager in managers:
            if manager.release(self.session_id, self.route_epoch):
                shards_released += 1
        buffers_released = 0
        for release in buffers:
            release()
            buffers_released += 1
        return CancellationOutcome(
            session_id=self.session_id,
            route_epoch=self.route_epoch,
            shards_released=shards_released,
            buffers_released=buffers_released,
        )


# --------------------------------------------------------------------------- #
# Idempotency: duplicate steps are no-ops; uncertain mutations never replay.
# --------------------------------------------------------------------------- #


class StepPhase(str, Enum):
    IN_FLIGHT = "in-flight"
    COMMITTED = "committed"
    UNCERTAIN = "uncertain"


class UncertainMutationError(FailureSemanticsError):
    """Raised when a caller tries to replay a step whose outcome is uncertain.

    A step is uncertain when its mutation may or may not have been applied (worker
    death / stream reset mid-append). Replaying it silently could double-apply KV
    or bill unverified work, so the ledger refuses: the caller must verify against
    the actual KV length or restart from token zero on a fresh epoch instead.
    """


@dataclass(frozen=True)
class StepKey:
    """Identity of one idempotent stream step within a route epoch."""

    session_id: str
    route_epoch: int
    step_index: int


@dataclass(frozen=True)
class StepDisposition:
    """What :meth:`IdempotencyLedger.begin` decided for a step."""

    fresh: bool
    token: int | None = None

    @property
    def duplicate(self) -> bool:
        return not self.fresh


class IdempotencyLedger:
    """Records committed/uncertain stream steps so duplicates never re-mutate.

    Keyed by ``(session, epoch, step_index)`` — the protocol's idempotency step.

    * :meth:`begin` on a *fresh* key marks it in-flight and returns "execute".
    * :meth:`begin` on a *committed* key returns the recorded token so a duplicate
      delivery is a no-op (idempotent replay).
    * :meth:`begin` on an *in-flight* or *uncertain* key raises
      :class:`UncertainMutationError` — a concurrent duplicate or a replay of an
      unverified mutation is never silently applied.
    """

    def __init__(self) -> None:
        self._phase: dict[StepKey, StepPhase] = {}
        self._token: dict[StepKey, int] = {}
        self._lock = threading.Lock()

    def begin(self, key: StepKey) -> StepDisposition:
        with self._lock:
            phase = self._phase.get(key)
            if phase is None:
                self._phase[key] = StepPhase.IN_FLIGHT
                return StepDisposition(fresh=True)
            if phase is StepPhase.COMMITTED:
                return StepDisposition(fresh=False, token=self._token[key])
            # IN_FLIGHT (concurrent duplicate) or UNCERTAIN (post-crash replay):
            # both are unverified and must not be silently re-applied.
            raise UncertainMutationError(
                f"step {key.step_index} for session {key.session_id[:8]} epoch "
                f"{key.route_epoch} is {phase.value}; refusing silent replay"
            )

    def commit(self, key: StepKey, token: int) -> None:
        with self._lock:
            self._phase[key] = StepPhase.COMMITTED
            self._token[key] = int(token)

    def mark_uncertain(self, key: StepKey, detail: str = "") -> None:
        with self._lock:
            # A committed step is verified; never downgrade it.
            if self._phase.get(key) is StepPhase.COMMITTED:
                return
            self._phase[key] = StepPhase.UNCERTAIN

    def phase_of(self, key: StepKey) -> StepPhase | None:
        with self._lock:
            return self._phase.get(key)

    def committed_token(self, key: StepKey) -> int | None:
        with self._lock:
            return self._token.get(key)

    def has_uncertain(self) -> bool:
        with self._lock:
            return any(p is StepPhase.UNCERTAIN for p in self._phase.values())


# --------------------------------------------------------------------------- #
# Restart / alpha failover: from token zero on a fresh compatible route.
# --------------------------------------------------------------------------- #


class RestartController:
    """Alpha failover that restarts from token zero, never importing prior KV.

    RALPH runtime decision #14: when the alpha (the head owning embedding + final
    head) fails, the route retries from token zero; unverified KV is never
    migrated. :meth:`failover` opens the *next* route epoch and releases every
    node-local shard's prior-epoch KV, so the restart begins with empty caches. The
    KV manager then treats the failed epoch as stale (a later reference to it is
    rejected), which is what keeps a half-computed cache from being reused.
    """

    def __init__(self, managers: Sequence[HotKvStateManager]) -> None:
        self._managers = list(managers)

    def failover(self, session_id: str, failed_epoch: int) -> int:
        """Advance to a fresh epoch and drop the failed epoch's KV on every shard."""
        new_epoch = int(failed_epoch) + 1
        for manager in self._managers:
            manager.release(session_id, failed_epoch)
        return new_epoch

    def assert_fresh_start(self, session_id: str, new_epoch: int) -> None:
        """Verify no shard carries KV for the new epoch (a true token-zero restart).

        Any residual KV under the new epoch would be unverified imported state;
        this fails closed so a restart can never silently attend over it.
        """
        for manager in self._managers:
            result = manager.resolve(session_id, new_epoch)
            if not isinstance(result, CacheMiss):
                raise FailureSemanticsError(
                    f"restart epoch {new_epoch} for session {session_id[:8]} is not "
                    "empty; refusing to import unverified KV"
                )


# --------------------------------------------------------------------------- #
# Billing / work records.
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class WorkRecord:
    """A typed unit of served work, distinguishing what may be billed.

    ``tokens`` counts only *committed* generated tokens. Only a
    :attr:`WorkStatus.COMPLETED` record is billable; cancelled/failed/unverified
    records carry their confirmed token count for observability but are excluded
    from billing so uncompleted or unverified work is never charged.
    """

    session_id: str
    route_epoch: int
    status: WorkStatus
    tokens: int
    failure_kind: FailureKind | None = None
    detail: str = ""

    @property
    def billable(self) -> bool:
        return self.status is WorkStatus.COMPLETED

    def to_dict(self) -> dict:
        return {
            "session_id": self.session_id,
            "route_epoch": self.route_epoch,
            "status": self.status.value,
            "tokens": self.tokens,
            "failure_kind": self.failure_kind.value if self.failure_kind else None,
            "detail": self.detail,
            "billable": self.billable,
        }


class WorkLedger:
    """Append-only ledger of :class:`WorkRecord`, split by billing status."""

    def __init__(self) -> None:
        self._records: list[WorkRecord] = []
        self._lock = threading.Lock()

    def record(self, record: WorkRecord) -> WorkRecord:
        with self._lock:
            self._records.append(record)
        return record

    def records(self) -> list[WorkRecord]:
        with self._lock:
            return list(self._records)

    def records_for(self, session_id: str) -> list[WorkRecord]:
        with self._lock:
            return [r for r in self._records if r.session_id == session_id]

    def billable_records(self) -> list[WorkRecord]:
        with self._lock:
            return [r for r in self._records if r.billable]

    def billable_tokens(self) -> int:
        """Total tokens that may be charged (completed work only)."""
        with self._lock:
            return sum(r.tokens for r in self._records if r.billable)

    def counts_by_status(self) -> dict[str, int]:
        counts: dict[str, int] = {s.value: 0 for s in WorkStatus}
        with self._lock:
            for record in self._records:
                counts[record.status.value] += 1
        return counts

    def to_dict(self) -> dict:
        with self._lock:
            records = [r.to_dict() for r in self._records]
        counts: dict[str, int] = {s.value: 0 for s in WorkStatus}
        for record in records:
            counts[record["status"]] += 1
        return {
            "schema_version": 1,
            "records": records,
            "counts_by_status": counts,
            "billable_tokens": sum(r["tokens"] for r in records if r["billable"]),
        }


# --------------------------------------------------------------------------- #
# The hardened single-session stream runner.
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class RunOutcome:
    """The typed result of one hardened generation attempt."""

    session_id: str
    route_epoch: int
    status: WorkStatus
    tokens: tuple[int, ...]
    failure_kind: FailureKind | None
    detail: str

    @property
    def completed(self) -> bool:
        return self.status is WorkStatus.COMPLETED

    @property
    def token_count(self) -> int:
        return len(self.tokens)

    @property
    def restartable(self) -> bool:
        return self.failure_kind in _RESTARTABLE

    def work_record(self) -> WorkRecord:
        return WorkRecord(
            session_id=self.session_id,
            route_epoch=self.route_epoch,
            status=self.status,
            tokens=len(self.tokens),
            failure_kind=self.failure_kind,
            detail=self.detail,
        )


@dataclass(frozen=True)
class FailoverResult:
    """The result of a run that may have restarted from token zero after a failure."""

    outcome: RunOutcome
    attempts: tuple[RunOutcome, ...]
    restarts: int

    @property
    def completed(self) -> bool:
        return self.outcome.completed

    def to_dict(self) -> dict:
        return {
            "final_status": self.outcome.status.value,
            "final_epoch": self.outcome.route_epoch,
            "restarts": self.restarts,
            "attempts": [
                {
                    "route_epoch": a.route_epoch,
                    "status": a.status.value,
                    "failure_kind": a.failure_kind.value if a.failure_kind else None,
                    "tokens": a.token_count,
                }
                for a in self.attempts
            ],
        }


class HardenedSessionRunner:
    """Drive one Route Session's decode stream with bounded failure semantics.

    The runner owns a single full-shard :class:`KvBoundaryAdapter` (head **and**
    tail, so a step samples a token) and threads every DGR-013 guarantee through a
    step loop:

    * every step is bounded by a :class:`DeadlineGuard` and can observe a
      :class:`CancellationToken`;
    * every step is idempotent through an :class:`IdempotencyLedger` (a duplicate
      returns the recorded token; an uncertain mutation is never replayed);
    * any failure releases this session's KV (cancellation) and is recorded as a
      typed :class:`WorkRecord` in the :class:`WorkLedger`;
    * :meth:`run_with_failover` restarts a transient failure from token zero on a
      fresh epoch via a :class:`RestartController`.
    """

    def __init__(
        self,
        adapter: KvBoundaryAdapter,
        *,
        clock: Callable[[], float] | None = None,
        work_ledger: WorkLedger | None = None,
        idempotency: IdempotencyLedger | None = None,
    ) -> None:
        if not (adapter.is_head and adapter.is_tail):
            raise FailureSemanticsError(
                "HardenedSessionRunner needs a full (head+tail) shard so decode "
                "steps sample tokens; got a partial range "
                f"(head={adapter.is_head} tail={adapter.is_tail})"
            )
        self._adapter = adapter
        self._manager: HotKvStateManager = adapter.manager
        self._clock = clock or time.monotonic
        self.work_ledger = work_ledger or WorkLedger()
        self.idempotency = idempotency or IdempotencyLedger()

    # -- single attempt -------------------------------------------------------

    def run(
        self,
        request: GenerationRequest,
        *,
        deadline: float | None = None,
        heartbeat_timeout: float | None = None,
        cancel_token: CancellationToken | None = None,
        heartbeat: Callable[[int], bool] | None = None,
        before_step: Callable[[int], None] | None = None,
    ) -> RunOutcome:
        """Run one attempt of ``request``; record and return a typed outcome.

        ``deadline`` (absolute, on the injected clock) and ``heartbeat_timeout``
        bound blocked steps. ``cancel_token`` lets a client cancel mid-stream.
        ``heartbeat(step)`` returns ``True`` when a heartbeat was heard before that
        step (resetting the health timer); ``before_step(step)`` is a fault-
        injection / clock-advance hook run before each step and may raise
        :class:`StreamTerminated` (e.g. a stream reset) or
        :class:`OperationCancelled`.
        """
        sid = request.session_id
        epoch = request.route_epoch
        guard = DeadlineGuard(
            deadline=deadline,
            heartbeat_timeout=heartbeat_timeout,
            clock=self._clock,
        )
        guard.start()
        tokens: list[int] = []
        current_key: StepKey | None = None
        try:
            # step 0 is the prefill (emits the first token); steps 1..N are decodes.
            for step_index in range(request.max_new_tokens):
                # before_step is the fault-injection / clock-advance hook and may
                # itself terminate the step (stream reset, cancel); run it first so
                # a fault it raises takes effect on this step, then re-check the
                # bounds it may have advanced (deadline / heartbeat / cancel).
                if before_step is not None:
                    before_step(step_index)
                if cancel_token is not None:
                    cancel_token.raise_if_cancelled()
                if heartbeat is not None and heartbeat(step_index):
                    guard.heartbeat()
                guard.check()

                current_key = StepKey(sid, epoch, step_index)
                disposition = self.idempotency.begin(current_key)
                if disposition.duplicate:
                    # Idempotent replay: reuse the recorded token, do not re-mutate.
                    assert disposition.token is not None
                    tokens.append(disposition.token)
                    continue

                token = self._execute_step(request, step_index, tokens)
                if isinstance(token, CacheMiss):
                    # The expected KV was gone; the append never started, so this is
                    # a certain (not uncertain) miss — restartable from token zero.
                    return self._finish_failure(
                        request,
                        tokens,
                        FailureKind.CACHE_MISS,
                        str(token),
                        cancel_token,
                    )
                self.idempotency.commit(current_key, token)
                tokens.append(token)
        except (StreamTerminated, OperationCancelled) as exc:
            return self._finish_failure(
                request, tokens, classify_exception(exc), str(exc), cancel_token
            )
        except (
            BoundaryContractError,
            StaleRouteEpochError,
            IncompatibleCacheRecipeError,
            KvCacheMissError,
        ) as exc:
            # Deterministic protocol/state errors, all validated before any KV
            # append committed — certain, not uncertain.
            return self._finish_failure(
                request, tokens, classify_exception(exc), str(exc), cancel_token
            )
        except UncertainMutationError as exc:
            # A replay of an unverified step reached the ledger — never silent.
            return self._finish_failure(
                request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token
            )
        except Exception as exc:  # noqa: BLE001 - unexpected == worker death
            # An unexpected error mid-step may have left the KV half-mutated; mark
            # the step uncertain so it can never be silently replayed, then fail
            # closed as unverified work.
            if current_key is not None:
                self.idempotency.mark_uncertain(current_key, str(exc))
            return self._finish_failure(
                request, tokens, FailureKind.WORKER_DEATH, str(exc), cancel_token
            )

        return self._finish_completed(request, tokens)

    def _execute_step(
        self, request: GenerationRequest, step_index: int, tokens: list[int]
    ) -> int | CacheMiss:
        sid = request.session_id
        epoch = request.route_epoch
        if step_index == 0:
            out = self._adapter.prefill(
                sid, epoch, token_ids=list(request.prompt_token_ids)
            )
        else:
            # expected_seq_len defends the KV layer against a desynchronised decode:
            # prompt positions plus the tokens already committed this run.
            expected = request.prompt_len + (step_index - 1)
            out = self._adapter.decode(
                sid,
                epoch,
                token_ids=[tokens[-1]],
                expected_seq_len=expected,
            )
        if isinstance(out, CacheMiss):
            return out
        if not isinstance(out, TailOutput):
            raise FailureSemanticsError(
                "full-shard step did not yield a sampled token; got "
                f"{type(out).__name__}"
            )
        return int(out.token_id)

    # -- failover across restarts --------------------------------------------

    def run_with_failover(
        self,
        request: GenerationRequest,
        controller: RestartController,
        *,
        max_restarts: int = 3,
        **run_kwargs: Any,
    ) -> FailoverResult:
        """Run ``request``, restarting a transient failure from token zero.

        On a restartable failure (worker death, stream reset, cache miss) the
        controller advances to a fresh epoch and drops the failed epoch's KV; the
        next attempt re-prefills the whole prompt from token zero. A deterministic
        failure (deadline, cancel, malformed bundle, stale epoch) is returned as-is
        — retrying it would hang or fail identically. Per-attempt fault-injection
        hooks (``before_step`` / ``heartbeat``) are only applied to the *first*
        attempt so a restart runs clean.
        """
        if max_restarts < 0:
            raise FailureSemanticsError("max_restarts must be >= 0")
        epoch = request.route_epoch
        attempts: list[RunOutcome] = []
        first_kwargs = run_kwargs
        for attempt in range(max_restarts + 1):
            attempt_request = replace(request, route_epoch=epoch)
            kwargs = first_kwargs if attempt == 0 else {}
            outcome = self.run(attempt_request, **kwargs)
            attempts.append(outcome)
            if outcome.completed or not outcome.restartable or attempt == max_restarts:
                return FailoverResult(
                    outcome=outcome, attempts=tuple(attempts), restarts=attempt
                )
            # Alpha failover: fresh epoch, drop prior-epoch KV on every shard, and
            # verify the new epoch starts empty (no unverified KV import).
            epoch = controller.failover(request.session_id, epoch)
            controller.assert_fresh_start(request.session_id, epoch)
        # Unreachable: the loop always returns, but keep the type-checker happy.
        raise FailureSemanticsError("run_with_failover exhausted without returning")

    # -- outcome bookkeeping --------------------------------------------------

    def _finish_completed(
        self, request: GenerationRequest, tokens: list[int]
    ) -> RunOutcome:
        outcome = RunOutcome(
            session_id=request.session_id,
            route_epoch=request.route_epoch,
            status=WorkStatus.COMPLETED,
            tokens=tuple(tokens),
            failure_kind=None,
            detail="",
        )
        self.work_ledger.record(outcome.work_record())
        return outcome

    def _finish_failure(
        self,
        request: GenerationRequest,
        tokens: list[int],
        kind: FailureKind,
        detail: str,
        cancel_token: CancellationToken | None,
    ) -> RunOutcome:
        # Cancellation semantics: release this session's local KV so a failed or
        # cancelled stream never leaks its cache. release() is idempotent.
        self._manager.release(request.session_id, request.route_epoch)
        if cancel_token is not None and kind is not FailureKind.CANCELLED:
            # Ensure downstream shards sharing the token also stop.
            cancel_token.cancel(kind.value)
        outcome = RunOutcome(
            session_id=request.session_id,
            route_epoch=request.route_epoch,
            status=work_status_for(kind),
            tokens=tuple(tokens),
            failure_kind=kind,
            detail=detail,
        )
        self.work_ledger.record(outcome.work_record())
        return outcome