ferrox-server 0.10.0

//! Shared token-generation loop for both the non-streaming and SSE
//! streaming `/v1/chat/completions` paths: sampling
//! (temperature/top-p/top-k/repetition penalty) with a greedy-argmax
//! path at `temperature<=0.0`,
//! plus stop-sequence handling that's correct even when a stop string
//! spans more than one generated token.

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use ferrox_core::cache::{KvBlockPool, KvCache, KvPoolExhausted as CacheKvPoolExhausted};
use ferrox_models::sampling::{Sampler, SamplingParams};
use ferrox_models::tokenizer::{prepend_bos, StopTokens};
use ferrox_models::{Ceiling, Decoder, Engine, KvElem, KvShape, PrefixCache, TextTokenizer};

use crate::budget::ContextCeiling;

use crate::json_mode::mask_logits_for_json;
use crate::model::ServerTokenizer;

#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
    #[error("prompt encoded to token id {token}, which is outside this model's vocabulary of {vocab_size} (its tokenizer does not match this checkpoint)")]
    TokenOutOfVocab { token: usize, vocab_size: usize },
    #[error("server is at capacity: the shared KV cache block pool has no free blocks for a new request; retry shortly")]
    KvPoolExhausted,
    /// The batch scheduler's admission queue is full. Distinct from
    /// `KvPoolExhausted`: nothing is exhausted, the server is simply
    /// further behind than it is willing to queue. Naming the depth and
    /// the cap keeps a retry storm diagnosable -- an operator can tell
    /// "too many clients" from "one request too big".
    #[error("server is at capacity: {queued} requests are already queued for the batch scheduler (limit {cap}); retry shortly")]
    QueueFull { queued: usize, cap: usize },
    /// The request cannot fit a ceiling that will not move: a typed
    /// refusal rather than an out-of-memory kill somewhere downstream.
    ///
    /// `binding` names *which* ceiling, because the two send an
    /// operator to different knobs -- `context_length_exceeded` is the
    /// request's size against what this deployment admits per request,
    /// `device_memory_budget_exceeded` is the whole server's KV budget.
    /// `estimated_bytes` and `limit_bytes` are the KV cost of the
    /// request and of the ceiling, so the arithmetic is checkable
    /// rather than asserted.
    ///
    /// Deliberately not a 503: an idle server refuses this identically,
    /// so "retry shortly" would be a lie.
    #[error("{binding}: {detail}")]
    KvBudgetExceeded {
        /// Machine-readable ceiling code from
        /// [`ferrox_models::Ceiling::code`].
        binding: &'static str,
        /// KV bytes this request would cost at its full length.
        estimated_bytes: u64,
        /// KV bytes the binding ceiling allows.
        limit_bytes: u64,
        /// Token positions the request asked for (prompt + max_tokens).
        positions: usize,
        /// Token positions the binding ceiling allows.
        positions_limit: usize,
        detail: String,
    },
}

impl DecodeError {
    /// Seconds to advise a client to wait before retrying, or `None`
    /// for an error retrying cannot fix.
    pub fn retry_after_secs(&self) -> Option<u64> {
        match self {
            DecodeError::TokenOutOfVocab { .. } => None,
            // Retrying an over-budget request changes nothing: the
            // ceiling it hit is the whole server, not the current load.
            DecodeError::KvBudgetExceeded { .. } => None,
            DecodeError::KvPoolExhausted | DecodeError::QueueFull { .. } => Some(1),
        }
    }
}

/// A shared `KvBlockPool` plus this server's admission-control wait
/// policy: how long a request is willing to retry acquiring its
/// per-layer caches before giving up, when the pool is momentarily
/// exhausted. `queue_wait: Duration::ZERO` (the default) means "try
/// once, reject immediately" -- the original reject-only behavior.
#[derive(Clone)]
pub struct KvPoolConfig {
    pub pool: Arc<Mutex<KvBlockPool>>,
    pub queue_wait: Duration,
}

/// Retries acquiring one `KvCache` per layer from `config.pool` until
/// either all of them succeed or `config.queue_wait` has elapsed since
/// the first attempt. Sleeping between attempts happens on whichever
/// thread calls this -- fine here since generation already runs on
/// tokio's blocking-thread pool (`spawn_blocking`), not an async
/// reactor thread that a `std::thread::sleep` would otherwise stall.
///
/// `max_seq_len` (this request's real worst-case sequence length --
/// prompt length plus `max_tokens`) is passed straight through to
/// `KvCache::with_pool`, so each layer's cache reserves enough blocks
/// for the *whole* request up front rather than growing mid-decode.
/// This isn't just an optimization: `Decoder::forward_token`/
/// `forward_batch` treat `KvCache::push` as infallible, so a pooled
/// cache that under-reserves at construction and then fails to
/// acquire another block later (because some other request took the
/// pool's remaining capacity in the meantime) would panic mid-decode
/// instead of failing this request cleanly at admission time -- caught
/// by a real panic during live testing before this fix.
/// The typed refusal for a request the KV pool can *never* satisfy, or
/// `None` when the pool could serve it once enough blocks come back.
///
/// A request holds one `KvCache` per layer and each reserves
/// `ceil(max_seq_len / block_size)` blocks up front (see
/// [`acquire_pooled_caches`]), so its whole-pool cost is
/// `n_layers * blocks_per_layer`. When that exceeds
/// `KvBlockPool::total_blocks` the request is refused by arithmetic
/// rather than by exhaustion: an *empty* pool would refuse it
/// identically, which is precisely the test for whether "retry shortly"
/// is a true statement. Before this existed, such a request slept
/// through `FERROX_KV_POOL_QUEUE_TIMEOUT_MS` and then got a 503 with a
/// `Retry-After` that could never come good.
///
/// Priced in real KV bytes through `ceiling`'s `KvShape` when one is
/// available, so the refusal names bytes rather than an opaque block
/// count. Positions, not bytes, decide it: the pool's own ledger is in
/// positions and this must agree with the acquisition it is predicting.
fn pool_immovable_refusal(
    decoder: &Decoder,
    config: &KvPoolConfig,
    max_seq_len: usize,
) -> Option<DecodeError> {
    let (block_size, total_blocks) = {
        let pool = config.pool.lock().unwrap_or_else(|p| p.into_inner());
        (pool.block_size(), pool.total_blocks())
    };
    if block_size == 0 || decoder.layers.is_empty() {
        return None;
    }
    let blocks_per_layer = max_seq_len.div_ceil(block_size).max(1);
    let needed = blocks_per_layer.saturating_mul(decoder.layers.len());
    if needed <= total_blocks {
        return None;
    }
    // The largest per-layer reservation the whole pool could cover, and
    // therefore the longest sequence it can ever hold. Floors to zero
    // for a pool too small for even one block per layer, which is an
    // honest answer: such a pool serves nothing.
    let blocks_per_layer_limit = total_blocks / decoder.layers.len();
    let positions_limit = blocks_per_layer_limit * block_size;
    let shape = KvShape::from_config(&decoder.config, KvElem::F32, 1);
    Some(DecodeError::KvBudgetExceeded {
        binding: Ceiling::DeviceMemory.code(),
        estimated_bytes: shape.kv_bytes_for_tokens(max_seq_len),
        limit_bytes: shape.kv_bytes_for_tokens(positions_limit),
        positions: max_seq_len,
        positions_limit,
        detail: format!(
            "request needs {needed} KV pool blocks ({max_seq_len} token positions at \
             {block_size} per block, across {} layers) but the whole pool is {total_blocks} \
             blocks; an idle server would refuse it identically",
            decoder.layers.len()
        ),
    })
}

fn acquire_pooled_caches(
    decoder: &Decoder,
    config: &KvPoolConfig,
    max_seq_len: usize,
) -> Result<Vec<KvCache>, CacheKvPoolExhausted> {
    let deadline = Instant::now() + config.queue_wait;
    loop {
        let attempt: Result<Vec<KvCache>, CacheKvPoolExhausted> = decoder
            .layers
            .iter()
            .map(|_| {
                KvCache::with_pool(
                    decoder.config.n_kv_heads,
                    decoder.config.head_dim,
                    Arc::clone(&config.pool),
                    max_seq_len,
                )
            })
            .collect();
        let now = Instant::now();
        if attempt.is_ok() || now >= deadline {
            return attempt;
        }
        std::thread::sleep(Duration::from_millis(10).min(deadline - now));
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
    Stop,
    Length,
    /// A canceller asked for this generation to stop; see the `cancel`
    /// module. Deliberately not folded into `Stop`: the tokens that did
    /// arrive are a partial answer, and a client that cannot tell a
    /// completed answer from an interrupted one will show the second as
    /// the first.
    Cancelled,
}

impl FinishReason {
    pub fn as_str(&self) -> &'static str {
        match self {
            FinishReason::Stop => "stop",
            FinishReason::Length => "length",
            // Not an OpenAI-defined value -- OpenAI has no cancel
            // endpoint to produce one. A client that does not know the
            // string still sees a terminated stream with *a* finish
            // reason, which is what stops it being read as truncation.
            FinishReason::Cancelled => "cancelled",
        }
    }
}

/// OpenAI-convention token accounting, reported in the response's
/// `usage` field. Counted from the exact token ids the generation loop
/// processed (prompt after BOS insertion, and every generated id), not
/// re-tokenized after the fact.
///
/// Defined in `ferrox-api` rather than here because the shape is part
/// of the public wire contract: the UI, `ferrox chat` and any external
/// client read these fields, so exactly one definition may exist (see
/// that crate's module docs for why the prefill and decode phases stay
/// separate).
pub use ferrox_api::Usage;

#[derive(Clone)]
pub struct GenerationParams {
    pub max_tokens: usize,
    pub sampling: SamplingParams,
    pub seed: u64,
    pub stop: Vec<String>,
    /// Stop strings that are exactly one token in this model's
    /// vocabulary, resolved once by whoever holds the tokenizer.
    ///
    /// Layer 1 of [`crate::stop`]: matched on the id, before the token
    /// is detokenized, so a control token stops the answer whatever it
    /// renders as. Empty when no stop string is a single token, or when
    /// the caller has no tokenizer to resolve them with.
    pub stop_token_ids: Vec<usize>,
    /// When true, constrain sampling toward JSON-safe token pieces and
    /// validate the emitted text is a JSON object (best-effort; see
    /// `json_mode` module).
    pub json_object: bool,
    /// Cooperative stop flag, polled once per decoded token.
    ///
    /// It rides on the params rather than being a separate argument
    /// because every path that already threads params through -- the
    /// streaming and non-streaming chat handlers, `/v1/completions`,
    /// the Anthropic surface -- then gets cancellation without a new
    /// parameter each, and a caller that has no cancellation to offer
    /// leaves it `None`. See the `cancel` module for the two tiers that
    /// set it.
    pub cancel: Option<crate::cancel::CancelToken>,
}

impl GenerationParams {
    /// Whether a canceller has asked this generation to stop.
    ///
    /// `None` -- no cancellation wired up -- is never cancelled, so a
    /// caller that does not care pays one branch and no atomic.
    pub(crate) fn is_cancelled(&self) -> bool {
        self.cancel.as_ref().is_some_and(|c| c.is_cancelled())
    }
}

fn chunked_prefill_tokens() -> Option<usize> {
    std::env::var("FERROX_CHUNKED_PREFILL")
        .ok()
        .and_then(|v| v.parse().ok())
        .filter(|&n| n > 0)
}

#[cfg(feature = "metal")]
fn cpu_kv_offload_enabled() -> bool {
    matches!(
        std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
        Some("1")
    )
}

/// Batched prompt prefill, optionally split into `FERROX_CHUNKED_PREFILL`-sized
/// chunks that append into the same KV caches.
fn forward_prompt_batch(
    decoder: &Decoder,
    tokens: &[usize],
    start_pos: usize,
    caches: &mut [KvCache],
) -> Vec<f32> {
    if let Some(chunk) = chunked_prefill_tokens() {
        let mut pos = start_pos;
        let mut last = Vec::new();
        for part in tokens.chunks(chunk) {
            last = decoder.forward_batch_last(part, pos, caches);
            pos += part.len();
        }
        last
    } else {
        decoder.forward_batch_last(tokens, start_pos, caches)
    }
}

/// Runs the prompt through `decoder`, then generates up to
/// `params.max_tokens` new tokens, calling `emit` with each newly-safe-
/// to-flush chunk of decoded text as it becomes available (see the
/// stop-sequence buffering note below). Returns the reason generation
/// stopped.
///
/// Stop-sequence correctness: a stop string can span more than one
/// generated token (e.g. stop=" END" while the tokenizer emits " "
/// and "END" as separate pieces), so text can't simply be flushed
/// token-by-token as soon as it's decoded -- the tail end of the
/// buffer might still turn into part of a stop match once the next
/// token arrives. This holds back the last `longest_stop_len - 1`
/// bytes of decoded text (respecting UTF-8 char boundaries) until
/// they're confirmed clean, the same buffering approach real
/// inference servers use for this exact reason.
#[allow(clippy::too_many_arguments)] // one clear parameter per concern; a
                                     // bundling struct here would just be GenerationParams's fields plus
                                     // decoder/tokenizer/stop_tokens/bos_id/prompt/kv_pool/prefix_cache/emit
                                     // re-wrapped for no real benefit at this call depth (two call sites,
                                     // both in this crate).
pub fn generate(
    decoder: &Decoder,
    tokenizer: &ServerTokenizer,
    stop_tokens: &StopTokens,
    bos_id: Option<usize>,
    prompt: &str,
    params: &GenerationParams,
    kv_pool: Option<&KvPoolConfig>,
    prefix_cache: Option<&Mutex<PrefixCache>>,
    ceiling: Option<&ContextCeiling>,
    mut emit: impl FnMut(&str),
) -> Result<(FinishReason, Usage), DecodeError> {
    let vocab_size = decoder.config.vocab_size;

    // Metal greedy GPU argmax: fold final_norm+lm_head+argmax into the
    // dense-stack CB and download one token id instead of hidden/vocab.
    // Thread-local so concurrent Arc<Decoder> requests do not race.
    #[cfg(feature = "metal")]
    let _metal_greedy_guard = {
        struct Guard;
        impl Drop for Guard {
            fn drop(&mut self) {
                ferrox_models::set_metal_greedy_argmax(false);
            }
        }
        if params.sampling.temperature <= 0.0 && ferrox_models::metal_greedy_gpu_enabled() {
            ferrox_models::set_metal_greedy_argmax(true);
            Some(Guard)
        } else {
            None
        }
    };

    let mut tokens = tokenizer.encode(prompt);
    prepend_bos(&mut tokens, bos_id);
    let prompt_tokens = tokens.len();
    if let Some(&bad) = tokens.iter().find(|&&t| t >= vocab_size) {
        return Err(DecodeError::TokenOutOfVocab {
            token: bad,
            vocab_size,
        });
    }

    // With a shared pool configured, admission control happens here:
    // each layer's cache reserves enough blocks up front for this
    // request's real worst case (prompt length + max_tokens), not just
    // one block -- see `acquire_pooled_caches`'s doc comment for why
    // under-reserving here would let a request panic mid-decode
    // instead of failing cleanly at admission time. If any layer can't
    // get that many blocks, `acquire_pooled_caches` retries (bounded by
    // `config.queue_wait`) before the request is rejected. Each failed
    // attempt's partial `Vec` is dropped immediately (via `collect`),
    // releasing any blocks it did acquire through `KvCache`'s `Drop`
    // impl before the next retry, so a request that ultimately gives up
    // leaves the pool exactly as it found it.
    //
    // Prefix-cache restoration only applies when there's no shared KV
    // block pool: a restored cache is a plain, unpooled clone (see
    // `KvCache`'s `Clone` doc comment), so combining it with pool-based
    // admission control would let a request's real memory usage
    // silently bypass the pool's bounded-memory guarantee. Not
    // supported together yet.
    let max_seq_len = tokens.len() + params.max_tokens;

    // The per-request context ceiling, checked before any KV is
    // acquired and before a single forward pass runs. This path used to
    // have no context ceiling at all: an oversized request either ran
    // until something else killed it, or -- with a pool configured --
    // waited out `queue_wait` and left with a 503 "retry shortly" about
    // a request an idle server would refuse identically. The ceiling is
    // the same `ContextCeiling` the batch scheduler admits on (see
    // `crate::budget`), so the two paths cannot disagree.
    if let Some(ceiling) = ceiling {
        if let Some(err) = ceiling.refusal(max_seq_len) {
            return Err(err);
        }
    }

    // A request whose worst case exceeds the *whole* pool is not a
    // request to retry: no amount of waiting frees blocks that do not
    // exist. Separated from `KvPoolExhausted` (503) for exactly the
    // reason the batch scheduler separates `kv_rejected_too_large` from
    // `queue_rejected` -- one says "come back later", the other says
    // "this will never work", and an operator sent to the wrong one
    // tunes the wrong knob.
    if let Some(config) = kv_pool {
        if let Some(err) = pool_immovable_refusal(decoder, config, max_seq_len) {
            return Err(err);
        }
    }

    let restored = if kv_pool.is_none() {
        prefix_cache.and_then(|pc| {
            let m = pc
                .lock()
                .unwrap_or_else(|p| p.into_inner())
                .find_longest_prefix(&tokens);
            (m.matched_len > 0).then_some(m)
        })
    } else {
        None
    };

    // Prompt tokens this request will *not* have to recompute. Reported
    // as `usage.cached_tokens` -- `Some(0)` when a prefix cache exists
    // and missed, `None` when there is no prefix cache to consult, so a
    // client can tell "no hit" from "no cache".
    let cached_tokens = restored
        .as_ref()
        .map(|m| m.matched_len)
        .or_else(|| (prefix_cache.is_some() && kv_pool.is_none()).then_some(0));

    let prefill_start = std::time::Instant::now();
    let mut pos;
    let mut logits: Vec<f32>;
    let mut caches: Vec<KvCache>;
    if let Some(m) = restored {
        caches = m
            .kv_caches
            .expect("matched_len > 0 always carries kv_caches");
        let suffix = &tokens[m.matched_len..];
        if suffix.is_empty() {
            if let Some(pl) = m.pending_logits {
                // The whole query was already processed and stored
                // verbatim before (e.g. an exact-repeat prompt with
                // unseeded sampling, so the whole-response cache
                // couldn't serve it) -- zero forward passes needed.
                pos = m.matched_len;
                logits = pl;
            } else {
                // Rare: our query exactly matches a strict PREFIX of a
                // longer stored entry (someone else's conversation
                // continued past this point), so there's no stored
                // "what comes next" for our shorter query. Back the
                // restored cache off by one position and reprocess
                // just that last matched token to get real logits,
                // rather than guessing.
                let back_to = m.matched_len - 1;
                for c in caches.iter_mut() {
                    c.truncate(back_to);
                }
                pos = back_to;
                logits = decoder.forward_token(tokens[back_to], pos, &mut caches);
                pos += 1;
            }
        } else {
            pos = m.matched_len;
            let mut l = Vec::new();
            for &tok in suffix {
                l = decoder.forward_token(tok, pos, &mut caches);
                pos += 1;
            }
            logits = l;
        }
    } else {
        caches = match kv_pool {
            Some(config) => acquire_pooled_caches(decoder, config, max_seq_len)
                .map_err(|_| DecodeError::KvPoolExhausted)?,
            None => decoder
                .layers
                .iter()
                .map(|_| KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim))
                .collect(),
        };
        // Process the prompt once, capturing the *last* call's logits
        // (which already predict the first generated token) instead of
        // discarding them. A prompt of zero tokens (bos_id unset and an
        // empty-string prompt encoding to nothing) has no real token to
        // seed from, so a synthetic id 0 bootstraps decoding.
        pos = 0;
        logits = if tokens.is_empty() {
            let l = decoder.forward_token(0, pos, &mut caches);
            pos += 1;
            l
        } else {
            // Batched prefill: one pass over the prompt (shared weight
            // traffic on CPU; fewer per-token bookkeeping costs). Last
            // row's logits predict the first generated token — same as
            // the sequential loop this replaces (see unit test below).
            // When `FERROX_CHUNKED_PREFILL` is set, split long prompts
            // into chunks that reuse the same KV caches.
            pos = tokens.len();
            forward_prompt_batch(decoder, &tokens, 0, &mut caches)
        };
    }

    let prefill_secs = prefill_start.elapsed().as_secs_f64();
    let decode_start = std::time::Instant::now();
    #[cfg(feature = "metal")]
    let kv_offload = cpu_kv_offload_enabled();
    let decode_token = |id: usize| tokenizer.decode(&[id]);

    // `logits` becomes the prediction for the position after each
    // generated token; every generated token gets exactly one
    // corresponding cache entry via the `step` closure below, matching
    // `prefix_cache`'s `pending_logits` expectations regardless of
    // whether the loop goes on to hit a stop sequence or max_tokens.
    // Time-to-first-token is stamped inside the `step` closure rather
    // than approximated as "prefill time": `step` is called immediately
    // after the first token is sampled, so this is the real moment the
    // user could have seen something. Threading it through the closure
    // instead of `sample_until_stop`'s signature keeps that function's
    // (already long) argument list unchanged, and the closure's mutable
    // borrow ends when it is dropped at the call's return.
    let mut first_token_at: Option<std::time::Instant> = None;
    let (finish, generated_ids, final_logits) = sample_until_stop(
        logits,
        pos,
        stop_tokens,
        params,
        |ids| tokenizer.decode(ids),
        |next, pos| {
            if first_token_at.is_none() {
                first_token_at = Some(std::time::Instant::now());
            }
            let l = decoder.forward_token(next, pos, &mut caches);
            #[cfg(feature = "metal")]
            if kv_offload {
                decoder.sync_metal_attn_kv_to_host(&mut caches);
            }
            l
        },
        &mut emit,
        if params.json_object {
            Some(&decode_token as &dyn Fn(usize) -> String)
        } else {
            None
        },
    );
    let decode_secs = decode_start.elapsed().as_secs_f64();
    logits = final_logits;
    let mut usage =
        Usage::new(prompt_tokens, generated_ids.len()).with_timings(prefill_secs, decode_secs);
    if let Some(at) = first_token_at {
        usage = usage.with_ttft(at.duration_since(prefill_start).as_secs_f64());
    }
    if let Some(cached) = cached_tokens {
        usage = usage.with_cached_tokens(cached);
    }

    // Store the full sequence this request actually processed (prompt
    // plus everything generated) so a future request sharing this
    // prefix -- the common multi-turn-chat case, where each turn's
    // prompt is the previous turn's full prompt+reply plus a little
    // more -- can skip recomputing it. `caches`/`logits` are exactly
    // in the right state for this: `logits` predicts whatever would
    // come after this sequence (the token that triggered an EOS/stop
    // match, if generation stopped that way, or the natural next
    // prediction if it ran to `max_tokens`), and every token in
    // `tokens`/`generated_ids` has exactly one corresponding cache
    // entry -- see the per-token push above. Skipped whenever a KV
    // pool is configured, for the same reason restoration is (see this
    // function's earlier comment).
    //
    // Metal dense-stack decode may leave host KvCache lagging the
    // Metal-resident KV; flush before storing so prefix restore gets
    // complete K/V.
    if kv_pool.is_none() {
        if let Some(pc) = prefix_cache {
            // Greedy Metal argmax returns a 1-element "logits" vec; that is
            // not a full pending distribution and must not be stored for
            // later (possibly non-greedy) prefix restores.
            if logits.len() == vocab_size {
                #[cfg(feature = "metal")]
                decoder.sync_metal_attn_kv_to_host(&mut caches);
                tokens.extend(generated_ids);
                pc.lock()
                    .unwrap_or_else(|p| p.into_inner())
                    .store(tokens, caches, logits);
            }
        }
    }

    Ok((finish, usage))
}

/// Shared sampling + stop-sequence-aware emission loop, given already-
/// primed `logits`/`pos` (the prompt has already been processed by the
/// caller) and a `step` closure that advances one position and returns
/// the new logits. Used by both `generate` (the GGUF path, whose own
/// KV-pool/prefix-cache-aware prompt priming stays separate above) and
/// `generate_engine` (any other `Engine`, with simpler priming and no
/// pooling/restoration) so the actual sampling/stop-sequence
/// correctness -- the part most worth not duplicating -- lives in
/// exactly one place. Returns the finish reason, the generated token
/// ids, and the final logits (the prediction for whatever would come
/// next), since `generate`'s prefix-cache storage needs both.
#[allow(clippy::too_many_arguments)] // one call site each from `generate`
                                     // and `generate_engine`; splitting it would only move the arguments.
fn sample_until_stop(
    mut logits: Vec<f32>,
    mut pos: usize,
    stop_tokens: &StopTokens,
    params: &GenerationParams,
    mut decode_one: impl FnMut(&[usize]) -> String,
    mut step: impl FnMut(usize, usize) -> Vec<f32>,
    mut emit: impl FnMut(&str),
    decode_token: Option<&dyn Fn(usize) -> String>,
) -> (FinishReason, Vec<usize>, Vec<f32>) {
    let mut matcher = crate::stop::StopMatcher::new(&params.stop, &params.stop_token_ids);
    let mut sampler = Sampler::new(params.seed);
    let mut generated_ids: Vec<usize> = Vec::with_capacity(params.max_tokens);
    let mut finish = FinishReason::Length;

    for _ in 0..params.max_tokens {
        // The one place cancellation is honoured, shared by `generate`
        // and `generate_engine`. Checked before sampling so a cancel
        // that lands between two tokens costs no further work, and
        // whatever `pending` already holds is still flushed below --
        // an interrupted answer keeps the tokens it earned.
        if params.is_cancelled() {
            finish = FinishReason::Cancelled;
            break;
        }
        let next = if params.json_object {
            if let Some(decode_token) = decode_token {
                let mut mask_fn = |scores: &mut [f32]| {
                    mask_logits_for_json(scores, decode_token);
                };
                sampler.sample_with_mask(
                    &logits,
                    &params.sampling,
                    &generated_ids,
                    Some(&mut mask_fn),
                )
            } else {
                sampler.sample(&logits, &params.sampling, &generated_ids)
            }
        } else {
            sampler.sample(&logits, &params.sampling, &generated_ids)
        };
        if stop_tokens.contains(next) {
            finish = FinishReason::Stop;
            break;
        }
        // Layer 1: before the token is detokenized or counted. A
        // control token the client asked to stop on is not part of the
        // answer, so it contributes neither an id nor a character --
        // exactly how `eos_id` is treated one line above.
        if matcher.is_stop_token(next) {
            finish = FinishReason::Stop;
            break;
        }
        generated_ids.push(next);
        logits = step(next, pos);
        pos += 1;

        // Layer 2: only text that can no longer become part of a stop
        // string leaves here.
        match matcher.push(&decode_one(&[next])) {
            crate::stop::StopStep::Emit(text) => {
                if !text.is_empty() {
                    emit(&text);
                }
            }
            crate::stop::StopStep::Matched(text) => {
                if !text.is_empty() {
                    emit(&text);
                }
                finish = FinishReason::Stop;
                break;
            }
        }
    }

    // Ended for some other reason (length, EOS, a cancel): whatever is
    // still withheld was output that no stop ever claimed.
    let tail = matcher.flush();
    if !tail.is_empty() {
        emit(&tail);
    }

    (finish, generated_ids, logits)
}

/// The generic counterpart to `generate`, for any `Engine` other than
/// `Decoder` (in practice, Kimi K3's `KimiEngine`). Deliberately
/// simpler: no KV block pool, no `PrefixCache` restoration -- Kimi's
/// KDA state is a fixed-size recurrent matrix that collapses history
/// irreversibly, so it can't support the truncate/restore operations
/// those features need (see `ferrox_models::engine`'s module docs). Every
/// request processes its full prompt from scratch against fresh
/// engine state.
pub fn generate_engine<E: Engine, T: TextTokenizer>(
    engine: &E,
    tokenizer: &T,
    stop_tokens: &StopTokens,
    bos_id: Option<usize>,
    prompt: &str,
    params: &GenerationParams,
    mut emit: impl FnMut(&str),
) -> Result<(FinishReason, Usage), DecodeError> {
    let vocab_size = engine.vocab_size();
    let mut tokens = tokenizer.encode(prompt);
    prepend_bos(&mut tokens, bos_id);
    let prompt_tokens = tokens.len();
    if let Some(&bad) = tokens.iter().find(|&&t| t >= vocab_size) {
        return Err(DecodeError::TokenOutOfVocab {
            token: bad,
            vocab_size,
        });
    }

    let mut state = engine.new_state();
    let mut pos = 0;
    let prefill_start = std::time::Instant::now();
    let logits = if tokens.is_empty() {
        let l = engine.forward_token(0, pos, &mut state);
        pos += 1;
        l
    } else {
        let mut l = Vec::new();
        for &tok in tokens.iter() {
            l = engine.forward_token(tok, pos, &mut state);
            pos += 1;
        }
        l
    };
    let prefill_secs = prefill_start.elapsed().as_secs_f64();
    let decode_start = std::time::Instant::now();

    // Timed for the same reason the `Decoder` path is: a UI that has to
    // wall-clock these engines instead cannot separate prefill from
    // decode, and Kimi/MLA prefill is sequential (one forward per prompt
    // token), so the two phases differ by more here, not less.
    let mut first_token_at: Option<std::time::Instant> = None;
    let (finish, generated_ids, _final_logits) = sample_until_stop(
        logits,
        pos,
        stop_tokens,
        params,
        |ids| tokenizer.decode(ids),
        |next, pos| {
            if first_token_at.is_none() {
                first_token_at = Some(std::time::Instant::now());
            }
            engine.forward_token(next, pos, &mut state)
        },
        &mut emit,
        None,
    );
    let decode_secs = decode_start.elapsed().as_secs_f64();

    let mut usage =
        Usage::new(prompt_tokens, generated_ids.len()).with_timings(prefill_secs, decode_secs);
    if let Some(at) = first_token_at {
        usage = usage.with_ttft(at.duration_since(prefill_start).as_secs_f64());
    }
    Ok((finish, usage))
}

/// The earliest byte offset in `text` at which any of `stops` begins,
/// or `None` if none match yet.
pub(crate) fn earliest_stop_match(text: &str, stops: &[String]) -> Option<usize> {
    stops
        .iter()
        .filter(|s| !s.is_empty())
        .filter_map(|s| text.find(s.as_str()))
        .min()
}

/// The largest char boundary `<= idx`. `str::floor_char_boundary` is
/// still nightly-only in stable Rust as of this writing; this is the
/// same walk-backward-to-a-boundary logic on stable.
pub(crate) fn floor_char_boundary(s: &str, idx: usize) -> usize {
    let mut i = idx.min(s.len());
    while i > 0 && !s.is_char_boundary(i) {
        i -= 1;
    }
    i
}

#[cfg(test)]
mod tests {
    use super::*;
    use ferrox_models::config::test_dense_fixture;

    fn small_decoder() -> Decoder {
        Decoder::new_random_small(test_dense_fixture(), 2, 256)
    }

    fn greedy_params(max_tokens: usize) -> GenerationParams {
        GenerationParams {
            max_tokens,
            sampling: SamplingParams::default(),
            seed: 1,
            stop: Vec::new(),
            stop_token_ids: Vec::new(),
            json_object: false,
            cancel: None,
        }
    }

    /// Regression test for a real bug caught by close reading, not by
    /// any earlier test (the earlier tests only checked `generate`
    /// against a test helper that replicated the same buggy pattern,
    /// so they never could have caught it): `generate`'s original
    /// prompt-processing loop pushed every prompt token into the KV
    /// cache once (correct), then its first generation-loop iteration
    /// re-processed the *last* prompt token a second time via another
    /// `forward_token` call at the wrong position (`tokens.len()`
    /// instead of its real position `tokens.len() - 1`) just to obtain
    /// logits -- silently duplicating that token in the cache with a
    /// different RoPE rotation applied, corrupting every subsequent
    /// position's attention. Fixed by capturing the prompt loop's own
    /// last-iteration logits instead of discarding and re-deriving
    /// them. This locks in the fixed pattern (which `generate` uses
    /// internally) against `forward_batch`'s independent ground truth.
    #[test]
    fn prompt_processing_matches_forward_batch_ground_truth_with_no_duplicate_position() {
        let decoder = small_decoder();
        let tokens = vec![1usize, 2, 3, 4];

        let mut fresh_caches: Vec<KvCache> = decoder
            .layers
            .iter()
            .map(|_| KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim))
            .collect();
        let batch_logits = decoder.forward_batch(&tokens, 0, &mut fresh_caches);
        let ground_truth_next_logits = batch_logits.last().unwrap().clone();

        // The exact pattern `generate` now uses: one forward_token call
        // per prompt token, keeping the last call's logits.
        let mut caches: Vec<KvCache> = decoder
            .layers
            .iter()
            .map(|_| KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim))
            .collect();
        let mut logits = Vec::new();
        for (pos, &tok) in tokens.iter().enumerate() {
            logits = decoder.forward_token(tok, pos, &mut caches);
        }

        assert_eq!(
            caches[0].seq_len, fresh_caches[0].seq_len,
            "must not push any position beyond the real prompt length"
        );
        // Tolerance, not bit equality. `forward_token` goes through
        // `WeightMatrix::apply` (one activation) and `forward_batch`
        // through `apply_batch`, and since the CPU batch GEMM kernels
        // landed those are different kernels with different f32
        // accumulation orders -- on aarch64 the batched path uses the
        // i8mm interleave-8 GEMM. So they agree to last-ulp, not
        // bit-for-bit. llama.cpp has the same batch-vs-sequential
        // property. The engine-side twin of this test was relaxed for
        // the same reason; this one only shows up on a host where the
        // aarch64 kernels actually run.
        assert_eq!(logits.len(), ground_truth_next_logits.len());
        for (i, (a, b)) in logits.iter().zip(&ground_truth_next_logits).enumerate() {
            assert!(
                (a - b).abs() <= 1e-5 * a.abs().max(1.0),
                "logit {i} predicting the first generated token: sequential {a} vs forward_batch {b}"
            );
        }
    }

    /// End-to-end version of the same property: `generate`'s full
    /// greedy decode loop (prompt processing + iterative generation)
    /// must produce exactly the token sequence an independent
    /// step-by-step computation (via `forward_batch` for the prompt,
    /// then `forward_token` once per new position, argmax at each
    /// step) would produce. Restricted to ASCII byte values so
    /// `ServerTokenizer::Byte`'s `decode` is lossless in both
    /// directions and the generated text can be compared back to
    /// token ids exactly.
    #[test]
    fn generate_greedy_output_matches_independent_step_by_step_computation() {
        let decoder = small_decoder();
        let prompt_ids = vec![1usize, 2, 3];
        let prompt = String::from_utf8(prompt_ids.iter().map(|&b| b as u8).collect()).unwrap();
        let max_tokens = 8;

        // Independent computation: forward_batch over the prompt, then
        // one forward_token + argmax per new position, decoding each
        // generated id one at a time and concatenating -- exactly
        // `generate`'s own decode granularity (`ServerTokenizer::Byte`
        // is lossy per non-ASCII byte, so decoding token-by-token vs.
        // decoding the whole sequence at once are not equivalent; this
        // must replicate the real call pattern, not just the ids).
        let mut caches: Vec<KvCache> = decoder
            .layers
            .iter()
            .map(|_| KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim))
            .collect();
        let mut logits = decoder
            .forward_batch(&prompt_ids, 0, &mut caches)
            .pop()
            .unwrap();
        let mut expected_text = String::new();
        for pos in (prompt_ids.len()..).take(max_tokens) {
            let next = logits
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
                .map(|(i, _)| i)
                .unwrap();
            expected_text.push_str(&ServerTokenizer::Byte.decode(&[next]));
            logits = decoder.forward_token(next, pos, &mut caches);
        }

        let mut actual_text = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(max_tokens),
            None,
            None,
            None,
            |s| actual_text.push_str(s),
        )
        .unwrap();

        assert_eq!(actual_text, expected_text);
    }

    #[test]
    fn rejects_out_of_vocab_prompt_tokens() {
        // ByteTokenizer emits raw bytes; vocab 32 makes ASCII letters OOV.
        let decoder = Decoder::new_random_small(test_dense_fixture(), 2, 32);
        let result = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            "hello",
            &greedy_params(4),
            None,
            None,
            None,
            |_| {},
        );
        assert!(matches!(result, Err(DecodeError::TokenOutOfVocab { .. })));
    }

    #[test]
    fn greedy_generation_hits_length_without_eos() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let mut chunks = String::new();
        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            None,
            None,
            |s| chunks.push_str(s),
        )
        .unwrap();
        assert_eq!(finish, FinishReason::Length);
    }

    /// A cancel raised while the loop is running must actually stop it,
    /// keep whatever was already decoded, and say `Cancelled` -- not
    /// `Stop`, which a client would render as a finished answer.
    #[test]
    fn a_cancelled_generation_stops_early_and_keeps_its_tokens() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let cancel = crate::cancel::CancelToken::new();

        let mut params = greedy_params(200);
        params.cancel = Some(cancel.clone());

        let mut chunks = String::new();
        let mut emitted = 0usize;
        let (finish, usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &params,
            None,
            None,
            None,
            |s| {
                chunks.push_str(s);
                emitted += 1;
                // Stands in for the socket dropping (or `/v1/cancel`
                // arriving) a few tokens into the answer.
                if emitted == 3 {
                    cancel.cancel();
                }
            },
        )
        .unwrap();

        assert_eq!(finish, FinishReason::Cancelled);
        assert!(
            usage.completion_tokens < 200,
            "cancelling did not shorten the decode: {} tokens",
            usage.completion_tokens
        );
        assert!(
            !chunks.is_empty(),
            "the tokens decoded before the cancel must survive it"
        );
    }

    /// A generation nobody cancelled must be untouched by the machinery
    /// -- the flag is polled every token, so a bug here would shorten
    /// every answer on the server.
    #[test]
    fn an_uncancelled_generation_runs_to_its_normal_end() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let mut params = greedy_params(5);
        params.cancel = Some(crate::cancel::CancelToken::new());

        let (finish, usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &params,
            None,
            None,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(finish, FinishReason::Length);
        assert_eq!(usage.completion_tokens, 5);
    }

    /// Discovers the real greedy-argmax next-token id after `prompt_ids`
    /// via `forward_batch` (ground truth: its last returned row predicts
    /// the token immediately after the full prompt, exactly what
    /// `generate` computes as its first generation-loop `logits` value)
    /// -- not by decoding it to text and reading bytes back, which is
    /// lossy for `ByteTokenizer`: a standalone byte >= 128 is not valid
    /// UTF-8 on its own, so `String::from_utf8_lossy` replaces it with
    /// the 3-byte U+FFFD replacement character, and reading
    /// `s.bytes().next()` off that recovers 0xEF (239), not the
    /// original token id.
    fn greedy_next_token_after(decoder: &Decoder, prompt_ids: &[usize]) -> usize {
        let mut caches: Vec<KvCache> = decoder
            .layers
            .iter()
            .map(|_| KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim))
            .collect();
        let logits = decoder
            .forward_batch(prompt_ids, 0, &mut caches)
            .pop()
            .unwrap();
        logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap()
    }

    #[test]
    fn eos_token_stops_generation_before_max_tokens() {
        let decoder = small_decoder();
        let prompt_ids = vec![1usize, 2];
        let prompt = String::from_utf8(prompt_ids.iter().map(|&b| b as u8).collect()).unwrap();

        // ByteTokenizer::encode is a lossless direct byte->id mapping
        // (only decode is lossy, see greedy_next_token_after's doc
        // comment), so `generate`'s internal prompt replay reaches
        // exactly the same state as this direct computation.
        let eos = greedy_next_token_after(&decoder, &prompt_ids);

        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::from_eos(Some(eos)),
            None,
            &prompt,
            &greedy_params(50),
            None,
            None,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(
            finish,
            FinishReason::Stop,
            "generation must stop as soon as the greedy-chosen token matches eos_id, not run to max_tokens"
        );
    }

    /// The bug this replaced: every server decode loop carried a single
    /// `eos_id`, so a Llama-3 checkpoint (whose `eos_token_id` is
    /// `<|end_of_text|>` while turns end with `<|eot_id|>`) or a gemma-2
    /// one (`<end_of_turn>`) ran past the end of its own turn to
    /// `max_tokens` over HTTP even after `eog_token_ids` landed for the
    /// CLI. Here the metadata EOS is deliberately a token the model will
    /// never pick, and the *turn ender* is the greedy next token: only a
    /// loop that consults the whole set stops.
    #[test]
    fn a_turn_ender_that_is_not_the_metadata_eos_still_stops_generation() {
        let decoder = small_decoder();
        let prompt_ids = vec![1usize, 2];
        let prompt = String::from_utf8(prompt_ids.iter().map(|&b| b as u8).collect()).unwrap();
        let turn_ender = greedy_next_token_after(&decoder, &prompt_ids);
        let never_sampled = (turn_ender + 1) % decoder.config.vocab_size;

        let stop = StopTokens::from_eos(Some(never_sampled)).with_id(Some(turn_ender));
        let (finish, usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &stop,
            None,
            &prompt,
            &greedy_params(50),
            None,
            None,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(finish, FinishReason::Stop);
        assert_eq!(
            usage.completion_tokens, 0,
            "the very first sampled token was the turn ender"
        );
    }

    #[test]
    fn a_stop_sequence_that_never_matches_does_not_drop_any_generated_content() {
        // The hold-back buffering (needed so a stop sequence spanning
        // more than one token is never partially flushed) must not
        // silently swallow output when no stop sequence ever matches:
        // the final emitted text must be byte-for-byte identical to an
        // otherwise-identical run with no stop sequences configured at
        // all, since the buffering is purely about *when* text is
        // flushed, never *whether* it is.
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8]).unwrap();

        let mut baseline = String::new();
        let (baseline_finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(20),
            None,
            None,
            None,
            |s| baseline.push_str(s),
        )
        .unwrap();

        let mut with_unmatchable_stop = String::new();
        let (stop_finish, _usage2) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &GenerationParams {
                max_tokens: 20,
                sampling: SamplingParams::default(),
                seed: 1,
                stop: vec!["ZZ_NEVER_MATCHES_ZZ".to_string()],
                stop_token_ids: Vec::new(),
                json_object: false,
                cancel: None,
            },
            None,
            None,
            None,
            |s| with_unmatchable_stop.push_str(s),
        )
        .unwrap();

        assert_eq!(baseline_finish, FinishReason::Length);
        assert_eq!(stop_finish, FinishReason::Length);
        assert_eq!(with_unmatchable_stop, baseline);
    }

    /// A decode loop whose next token is scripted, so the two stop
    /// layers can be exercised on exactly the token sequence they are
    /// meant to react to rather than on whatever a random decoder
    /// happens to emit.
    ///
    /// `script` is the token id produced at each step; `render` is how
    /// each id detokenizes.
    fn run_scripted(
        script: &[usize],
        render: impl Fn(usize) -> String,
        params: &GenerationParams,
    ) -> (FinishReason, Vec<usize>, Vec<String>) {
        let vocab = script.iter().copied().max().unwrap_or(0) + 2;
        let logits_for = |id: usize| {
            let mut v = vec![0.0f32; vocab];
            v[id] = 10.0;
            v
        };
        let mut next = 0usize;
        let mut take = || {
            let id = script
                .get(next)
                .copied()
                .unwrap_or(script[script.len() - 1]);
            next += 1;
            id
        };
        let first = logits_for(take());
        let mut chunks: Vec<String> = Vec::new();
        let (finish, ids, _) = sample_until_stop(
            first,
            0,
            &StopTokens::from_eos(None),
            params,
            |ids| ids.iter().copied().map(&render).collect::<String>(),
            |_tok, _pos| logits_for(take()),
            |chunk| chunks.push(chunk.to_string()),
            None,
        );
        (finish, ids, chunks)
    }

    fn scripted_params(max_tokens: usize) -> GenerationParams {
        GenerationParams {
            max_tokens,
            sampling: SamplingParams {
                temperature: 0.0,
                ..SamplingParams::default()
            },
            seed: 1,
            stop: Vec::new(),
            stop_token_ids: Vec::new(),
            json_object: false,
            cancel: None,
        }
    }

    /// Layer 1: a stop *token* ends the answer, and the token itself
    /// never appears in it -- the same treatment EOS already gets,
    /// which is the point. A control token the client named is no more
    /// part of the output than the end-of-sequence token is.
    ///
    /// Confirmed to FAIL (runs to all 6 tokens, emitting "aabaab") when
    /// the `is_stop_token` check is removed from `sample_until_stop`.
    #[test]
    fn a_token_level_stop_ends_generation_and_never_reaches_the_output() {
        let render = |id: usize| char::from(b'a' + id as u8).to_string();
        let script = [0usize, 0, 1, 0, 0, 1];

        let (finish, ids, chunks) = run_scripted(&script, render, &scripted_params(6));
        assert_eq!(finish, FinishReason::Length);
        assert_eq!(ids, script.to_vec());
        assert_eq!(chunks.concat(), "aabaab");

        let (finish, ids, chunks) = run_scripted(
            &script,
            render,
            &GenerationParams {
                stop_token_ids: vec![1],
                ..scripted_params(6)
            },
        );
        assert_eq!(finish, FinishReason::Stop);
        assert_eq!(ids, vec![0, 0], "the stop token is not part of the answer");
        assert_eq!(
            chunks.concat(),
            "aa",
            "the stop token must not be rendered into the output"
        );
    }

    /// The case layer 2 provably cannot cover: a control token that
    /// renders as the empty string. The text scan has nothing to match
    /// on, so only matching the id ends the answer.
    #[test]
    fn a_stop_token_that_renders_as_nothing_is_still_a_stop() {
        // Token 1 detokenizes to "", as real special tokens can.
        let render = |id: usize| {
            if id == 1 {
                String::new()
            } else {
                char::from(b'a' + id as u8).to_string()
            }
        };
        let script = [0usize, 1, 0, 0];

        // The text layer alone: the stop string never appears, so
        // generation runs to its limit.
        let (finish, _, chunks) = run_scripted(
            &script,
            render,
            &GenerationParams {
                stop: vec!["<|end|>".to_string()],
                ..scripted_params(4)
            },
        );
        assert_eq!(finish, FinishReason::Length);
        assert_eq!(chunks.concat(), "aaa");

        // With the id resolved, it stops where it should.
        let (finish, ids, chunks) = run_scripted(
            &script,
            render,
            &GenerationParams {
                stop: vec!["<|end|>".to_string()],
                stop_token_ids: vec![1],
                ..scripted_params(4)
            },
        );
        assert_eq!(finish, FinishReason::Stop);
        assert_eq!(ids, vec![0]);
        assert_eq!(chunks.concat(), "a");
    }

    /// Layer 2: nothing that turns out to be part of the stop string
    /// is ever emitted. `emit` is append-only -- SSE has no way to take
    /// a chunk back -- so over-emitting is not a display glitch, it is
    /// the stop sequence failing to do the one thing it promises.
    ///
    /// The script spells "ab" (a partial match that is disproved) and
    /// then "abc" (the real one), so the buffer has to hold, release,
    /// and hold again.
    #[test]
    fn nothing_that_becomes_part_of_the_stop_is_ever_emitted() {
        let render = |id: usize| char::from(b'a' + id as u8).to_string();
        let script = [0usize, 1, 0, 1, 2, 0];
        let params = GenerationParams {
            stop: vec!["abc".to_string()],
            ..scripted_params(6)
        };

        let (finish, _, chunks) = run_scripted(&script, render, &params);
        assert_eq!(finish, FinishReason::Stop);
        assert_eq!(
            chunks.concat(),
            "ab",
            "the answer is everything before the stop, and nothing after it"
        );

        // Append-only means every intermediate state must already be a
        // prefix of that: emitting one character too many can never be
        // undone.
        let mut seen = String::new();
        for chunk in &chunks {
            seen.push_str(chunk);
            assert!(
                "ab".starts_with(&seen),
                "the stream ran ahead of the answer: {seen:?} (chunks: {chunks:?})"
            );
        }
    }

    /// A partial match that is disproved is released with the very
    /// token that disproves it, not carried to the end of the answer.
    ///
    /// The conservative alternative -- always withhold
    /// `longest_stop - 1` bytes -- is equally safe and permanently
    /// leaves the stream that many bytes behind the model, for a match
    /// that in most chunks is not even beginning.
    ///
    /// Confirmed to FAIL (first chunk is "a", not "abd") when
    /// `partial_suffix_len` is replaced by that fixed hold-back.
    #[test]
    fn a_disproved_partial_is_released_by_the_token_that_disproves_it() {
        let render = |id: usize| char::from(b'a' + id as u8).to_string();
        // "a", "b" are held as a possible "abc"; "d" settles it.
        let script = [0usize, 1, 3, 0];
        let (_, _, chunks) = run_scripted(
            &script,
            render,
            &GenerationParams {
                stop: vec!["abc".to_string()],
                ..scripted_params(4)
            },
        );
        assert_eq!(chunks.concat(), "abda", "no output is lost");
        assert_eq!(
            chunks.first().map(String::as_str),
            Some("abd"),
            "the whole disproved partial goes out at once: {chunks:?}"
        );
    }

    #[test]
    fn a_stop_sequence_that_does_match_truncates_output_before_it() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8]).unwrap();

        // Discover what greedy decode actually produces, then use a
        // substring of it (starting after the first character, so at
        // least one character of real output precedes the match) as a
        // stop sequence guaranteed to match.
        let mut baseline = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(20),
            None,
            None,
            None,
            |s| baseline.push_str(s),
        )
        .unwrap();
        let Some((cut, _)) = baseline.char_indices().nth(1) else {
            // Degenerate case for this decoder/seed: fewer than 2 chars
            // generated: nothing meaningful to truncate, skip.
            return;
        };
        let stop_str = baseline[cut..].to_string();
        if stop_str.is_empty() {
            return;
        }

        let mut truncated = String::new();
        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &GenerationParams {
                max_tokens: 20,
                sampling: SamplingParams::default(),
                seed: 1,
                stop: vec![stop_str],
                stop_token_ids: Vec::new(),
                json_object: false,
                cancel: None,
            },
            None,
            None,
            None,
            |s| truncated.push_str(s),
        )
        .unwrap();

        assert_eq!(finish, FinishReason::Stop);
        assert_eq!(truncated, baseline[..cut]);
    }

    #[test]
    fn usage_reports_both_phases_and_a_time_to_first_token() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let (_finish, usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            None,
            None,
            |_| {},
        )
        .unwrap();

        assert_eq!(usage.prompt_tokens, 3);
        assert_eq!(usage.completion_tokens, 5);
        let prefill = usage.prompt_eval_duration_ms.expect("prefill timed");
        let decode = usage.generation_duration_ms.expect("decode timed");
        let ttft = usage.time_to_first_token_ms.expect("first token timed");
        // TTFT is measured from the start of prefill, so it can never be
        // shorter than prefill, and the first of five tokens must land
        // before the decode loop finishes all five.
        assert!(ttft >= prefill, "ttft {ttft} < prefill {prefill}");
        assert!(
            ttft <= prefill + decode + 1.0,
            "ttft {ttft} exceeds the whole request"
        );
    }

    #[test]
    fn cached_tokens_distinguishes_a_miss_from_an_absent_prefix_cache() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();

        let (_f, no_cache) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(2),
            None,
            None,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(no_cache.cached_tokens, None, "no prefix cache configured");

        let pc = Mutex::new(PrefixCache::new(4));
        let (_f, miss) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(2),
            None,
            Some(&pc),
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(miss.cached_tokens, Some(0), "cache consulted, missed");

        // Second turn extends the first: three prompt tokens are reused.
        let longer = String::from_utf8(vec![1u8, 2, 3, 9]).unwrap();
        let (_f, hit) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &longer,
            &greedy_params(2),
            None,
            Some(&pc),
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(hit.cached_tokens, Some(3));
    }

    #[test]
    fn prefix_cache_reuses_a_shared_prefix_and_produces_the_same_output_as_a_fresh_run() {
        let decoder = small_decoder();
        let prompt1 = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let pc = Mutex::new(PrefixCache::new(4));

        let mut out1 = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt1,
            &greedy_params(5),
            None,
            Some(&pc),
            None,
            |s| out1.push_str(s),
        )
        .unwrap();
        assert_eq!(pc.lock().unwrap().stats().misses, 1);

        // prompt2's tokens (raw bytes, via ByteTokenizer) start with
        // prompt1's exact bytes -- the common multi-turn-chat shape.
        let prompt2 = String::from_utf8(vec![1u8, 2, 3, 9, 9]).unwrap();

        let mut out2_with_cache = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt2,
            &greedy_params(5),
            None,
            Some(&pc),
            None,
            |s| out2_with_cache.push_str(s),
        )
        .unwrap();
        let stats = pc.lock().unwrap().stats();
        assert_eq!(stats.hits, 1, "prompt2 must hit the stored prompt1 entry");
        assert_eq!(stats.total_positions_reused, 3);

        let mut out2_fresh = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt2,
            &greedy_params(5),
            None,
            None,
            None,
            |s| out2_fresh.push_str(s),
        )
        .unwrap();

        assert_eq!(
            out2_with_cache, out2_fresh,
            "restoring from the prefix cache must produce identical output to processing the whole prompt from scratch"
        );
    }

    #[test]
    fn prefix_cache_exact_repeat_skips_prompt_processing_via_pending_logits() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let pc = Mutex::new(PrefixCache::new(4));

        let mut out1 = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            Some(&pc),
            None,
            |s| out1.push_str(s),
        )
        .unwrap();

        // The exact same prompt again: the stored entry's `tokens` (the
        // full prompt+completion from the first call) starts with this
        // exact prompt, so this only matches the prompt-length prefix
        // of a longer stored entry, not an exact full-entry match --
        // covering the "no pending_logits available" fallback path,
        // not the zero-forward-pass shortcut. Still must produce
        // identical output to a from-scratch run.
        let mut out2_with_cache = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            Some(&pc),
            None,
            |s| out2_with_cache.push_str(s),
        )
        .unwrap();

        let mut out2_fresh = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            None,
            None,
            |s| out2_fresh.push_str(s),
        )
        .unwrap();

        assert_eq!(out2_with_cache, out2_fresh);
    }

    #[test]
    fn prefix_cache_is_not_consulted_when_a_kv_pool_is_configured() {
        let decoder = small_decoder(); // 2 layers
        let prompt = String::from_utf8(vec![1u8, 2, 3]).unwrap();
        let pc = Mutex::new(PrefixCache::new(4));
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 2)));
        let config = pool_config(pool, Duration::ZERO);

        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            Some(&pc),
            None,
            |_| {},
        )
        .unwrap();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            Some(&pc),
            None,
            |_| {},
        )
        .unwrap();

        let stats = pc.lock().unwrap().stats();
        assert_eq!(
            stats.hits + stats.misses,
            0,
            "prefix cache must never be consulted while a KV pool is configured"
        );
    }

    fn pool_config(pool: Arc<Mutex<KvBlockPool>>, queue_wait: Duration) -> KvPoolConfig {
        KvPoolConfig { pool, queue_wait }
    }

    #[test]
    fn generate_succeeds_with_a_pool_that_has_enough_blocks() {
        let decoder = small_decoder(); // 2 layers
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 2)));
        let config = pool_config(pool.clone(), Duration::ZERO);

        let mut out = String::new();
        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            None,
            None,
            |s| out.push_str(s),
        )
        .unwrap();
        assert_eq!(finish, FinishReason::Length);
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            2,
            "every acquired block must be released once the request finishes"
        );
    }

    /// Regression test for a real bug caught by live testing (not by
    /// any unit test): with a small `block_size`, a request whose
    /// prompt + max_tokens exceeds one block used to reserve only one
    /// block per layer at admission time, then panic deep inside
    /// `Decoder::forward_token` once decode outgrew that block and the
    /// pool had nothing left to grow into (`KvCache::push` returning
    /// `Err` where `forward_token` assumes it can't fail). Fixed by
    /// having `acquire_pooled_caches` reserve blocks for the whole
    /// worst-case sequence length up front. This test must not panic --
    /// it must either succeed cleanly or fail at admission with
    /// `KvPoolExhausted`, never partway through decode.
    #[test]
    fn generate_reserves_enough_blocks_up_front_for_a_sequence_spanning_multiple_blocks() {
        let decoder = small_decoder(); // 2 layers
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap(); // 2 tokens via ByteTokenizer
        let max_tokens = 10;
        let block_size = 2;
        // prompt (2) + max_tokens (10) = 12 positions -> 6 blocks/layer * 2 layers = 12 blocks.
        let pool = Arc::new(Mutex::new(KvBlockPool::new(block_size, 12)));
        let config = pool_config(pool.clone(), Duration::ZERO);

        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(max_tokens),
            Some(&config),
            None,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(finish, FinishReason::Length);
        assert_eq!(pool.lock().unwrap().free_blocks(), 12);
    }

    /// One block short of the worst case means the pool can *never*
    /// serve this request, so the refusal is the immovable one -- a 400
    /// naming `device_memory_budget_exceeded`, not a 503 inviting a
    /// retry that an empty pool would refuse identically.
    ///
    /// Confirmed to FAIL when `pool_immovable_refusal` is removed from
    /// `generate` (the request falls through to the acquisition and
    /// comes back as the retryable `KvPoolExhausted`), and when its
    /// `needed <= total_blocks` comparison is loosened to `<`.
    #[test]
    fn generate_fails_at_admission_not_mid_decode_when_the_pool_cannot_cover_the_worst_case() {
        let decoder = small_decoder(); // 2 layers
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let max_tokens = 10;
        let block_size = 2;
        // One block short of the 12 the worst case (see the test
        // above) actually needs.
        let pool = Arc::new(Mutex::new(KvBlockPool::new(block_size, 11)));
        let config = pool_config(pool.clone(), Duration::ZERO);

        let result = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(max_tokens),
            Some(&config),
            None,
            None,
            |_| {},
        );
        let err = result.expect_err("11 blocks cannot cover a 12-block worst case");
        assert!(
            matches!(
                &err,
                DecodeError::KvBudgetExceeded { binding, positions, .. }
                    if *binding == ferrox_models::Ceiling::DeviceMemory.code()
                        && *positions == 12
            ),
            "expected an immovable device-memory refusal, got {err:?}"
        );
        assert_eq!(
            err.retry_after_secs(),
            None,
            "no wait frees blocks that do not exist"
        );
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            11,
            "a rejected request must leave the pool exactly as it found it"
        );
    }

    /// A one-block pool cannot hold a two-layer model's caches under any
    /// schedule, so this is the immovable refusal too.
    #[test]
    fn generate_rejects_the_request_without_leaking_blocks_when_the_pool_is_too_small() {
        let decoder = small_decoder(); // 2 layers -> needs 2 blocks, one per layer's cache
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 1)));
        let config = pool_config(pool.clone(), Duration::ZERO);

        let result = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            None,
            None,
            |_| {},
        );
        let err = result.expect_err("one block cannot hold two layers' caches");
        assert!(
            matches!(
                &err,
                DecodeError::KvBudgetExceeded { binding, .. }
                    if *binding == ferrox_models::Ceiling::DeviceMemory.code()
            ),
            "expected an immovable device-memory refusal, got {err:?}"
        );
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            1,
            "a rejected request must leave the pool exactly as it found it"
        );
    }

    #[test]
    fn generate_releases_blocks_so_back_to_back_requests_do_not_starve_the_pool() {
        let decoder = small_decoder(); // 2 layers
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        // Just enough for one request's caches at a time -- a second,
        // concurrent request would be rejected, but a *sequential*
        // second request must succeed once the first has returned its
        // blocks.
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 2)));
        let config = pool_config(pool.clone(), Duration::ZERO);

        for _ in 0..3 {
            let (finish, _usage) = generate(
                &decoder,
                &ServerTokenizer::Byte,
                &StopTokens::default(),
                None,
                &prompt,
                &greedy_params(5),
                Some(&config),
                None,
                None,
                |_| {},
            )
            .unwrap();
            assert_eq!(finish, FinishReason::Length);
        }
        assert_eq!(pool.lock().unwrap().free_blocks(), 2);
    }

    /// `queue_wait = 0` must reject on the first failed attempt rather
    /// than retry.
    ///
    /// The pool here is big enough for the request and *momentarily*
    /// held by someone else, which is the only situation in which
    /// `KvPoolExhausted` is an honest answer: a pool permanently too
    /// small is refused earlier, by `pool_immovable_refusal`, and would
    /// make the timing assertion below vacuous.
    #[test]
    fn generate_with_zero_queue_wait_rejects_immediately() {
        let decoder = small_decoder(); // 2 layers -> needs 2 blocks
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 2)));

        // Another in-flight request holding both blocks for longer than
        // this request is willing to wait for them.
        let holder_pool = pool.clone();
        let holder = std::thread::spawn(move || {
            let mut held = KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
            held.push(&[0.0], &[0.0]).unwrap(); // crosses into the second block
            std::thread::sleep(Duration::from_millis(200));
            drop(held);
        });
        std::thread::sleep(Duration::from_millis(15));

        let config = pool_config(pool, Duration::ZERO);
        let started = Instant::now();
        let result = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            None,
            None,
            |_| {},
        );
        assert!(
            matches!(result, Err(DecodeError::KvPoolExhausted)),
            "a pool that could serve this request once its holder lets go is momentary \
             exhaustion, which is retryable"
        );
        assert!(
            started.elapsed() < Duration::from_millis(50),
            "queue_wait=0 must reject on the first attempt, not retry: took {:?}",
            started.elapsed()
        );
        holder.join().unwrap();
    }

    /// The private path's context ceiling, and the property that makes
    /// it worth having: the refusal happens before any KV is acquired
    /// and before a single forward pass runs.
    ///
    /// Confirmed to FAIL when the `ceiling.refusal` block is removed
    /// from `generate` -- the request then runs to completion and
    /// returns `Ok`.
    #[test]
    fn a_request_past_the_context_ceiling_is_refused_before_any_kv_is_acquired() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 64)));
        let config = pool_config(pool.clone(), Duration::ZERO);
        let shape = KvShape::from_config(&decoder.config, KvElem::F32, 1);
        // Prompt is 2 tokens, max_tokens is 5, so the request asks for
        // 7 positions against a ceiling of 4.
        let ceiling = ContextCeiling::new(Some(4), shape);

        let err = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            None,
            Some(&ceiling),
            |_| panic!("no token may be emitted by a refused request"),
        )
        .expect_err("7 positions must not be admitted under a 4-position ceiling");
        match &err {
            DecodeError::KvBudgetExceeded {
                binding,
                positions,
                positions_limit,
                estimated_bytes,
                limit_bytes,
                ..
            } => {
                assert_eq!(*binding, ferrox_models::Ceiling::ContextLength.code());
                assert_eq!(*positions, 7);
                assert_eq!(*positions_limit, 4);
                assert_eq!(*estimated_bytes, shape.kv_bytes_for_tokens(7));
                assert_eq!(*limit_bytes, shape.kv_bytes_for_tokens(4));
            }
            other => panic!("expected a context-length refusal, got {other:?}"),
        }
        assert_eq!(err.retry_after_secs(), None, "a 400, not a retryable 503");
        assert_eq!(
            pool.lock().unwrap().free_blocks(),
            64,
            "the refusal must land before any block is taken"
        );
        assert_eq!(ceiling.refused(), 1);
    }

    /// A request that fits the ceiling is untouched by it: the same
    /// request that runs without a ceiling runs with one.
    ///
    /// Without this, a ceiling that refused everything would still pass
    /// the test above.
    #[test]
    fn a_request_inside_the_ceiling_is_admitted_unchanged() {
        let decoder = small_decoder();
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let shape = KvShape::from_config(&decoder.config, KvElem::F32, 1);
        let ceiling = ContextCeiling::new(Some(7), shape);

        let mut with = String::new();
        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            None,
            Some(&ceiling),
            |s| with.push_str(s),
        )
        .expect("7 positions fits a 7-position ceiling exactly");
        assert_eq!(finish, FinishReason::Length);

        let mut without = String::new();
        generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            None,
            None,
            None,
            |s| without.push_str(s),
        )
        .unwrap();
        assert_eq!(with, without, "an unbinding ceiling must change nothing");
        assert_eq!(ceiling.refused(), 0);
    }

    #[test]
    fn generate_with_a_queue_wait_succeeds_once_another_holder_releases_its_blocks() {
        let decoder = small_decoder(); // 2 layers, needs 2 blocks
        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
        let pool = Arc::new(Mutex::new(KvBlockPool::new(64, 2)));

        // Hold both blocks on another thread for a short while, then
        // release them -- simulating another in-flight request that's
        // about to finish.
        let holder_pool = pool.clone();
        let holder = std::thread::spawn(move || {
            let mut held = KvCache::with_pool(1, 1, holder_pool.clone(), 0).unwrap();
            held.push(&[0.0], &[0.0]).unwrap(); // crosses into needing the second block
            std::thread::sleep(Duration::from_millis(80));
            drop(held); // returns both blocks to the pool
        });
        // Give the holder a moment to actually acquire before we try.
        std::thread::sleep(Duration::from_millis(15));

        let config = pool_config(pool.clone(), Duration::from_millis(500));
        let (finish, _usage) = generate(
            &decoder,
            &ServerTokenizer::Byte,
            &StopTokens::default(),
            None,
            &prompt,
            &greedy_params(5),
            Some(&config),
            None,
            None,
            |_| {},
        )
        .unwrap();
        assert_eq!(
            finish,
            FinishReason::Length,
            "a sufficiently long queue_wait must let the request succeed once the holder releases"
        );
        holder.join().unwrap();
        assert_eq!(pool.lock().unwrap().free_blocks(), 2);
    }

    #[test]
    fn earliest_stop_match_finds_the_leftmost_match_across_multiple_stops() {
        assert_eq!(
            earliest_stop_match("hello world", &["world".to_string(), "hello".to_string()]),
            Some(0)
        );
        assert_eq!(
            earliest_stop_match("hello world", &["nope".to_string()]),
            None
        );
    }
}