car_inference/tasks/

generate.rs

//! Text generation with sampling.

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
use candle_core::Tensor;
use serde::{Deserialize, Serialize};

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
use crate::backend::CandleBackend;
#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
use crate::InferenceError;

/// How latency-sensitive this generation request is.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum RoutingWorkload {
    /// User-facing, interactive request where latency matters.
    #[default]
    Interactive,
    /// Batch job where latency matters somewhat, but quality/cost matter more.
    Batch,
    /// Background or offline work where latency is a weak concern.
    Background,
    /// Caller explicitly prefers on-device models. Distinct from
    /// `Background` (which is "this is a background job, latency
    /// barely matters"). The caller may be doing latency-sensitive
    /// interactive work but wants the privacy / cost / offline
    /// properties of local inference. Same `local_bonus` as
    /// `Background` plus a slightly more quality-aware weight profile
    /// — the caller chose local for a reason, not because the work is
    /// throwaway.
    LocalPreferred,
    /// Aggressive latency bias for time-to-first-token. Voice turns
    /// (specifically the fast track in the two-track sidecar pattern)
    /// pick this. Quality and cost are heavily downweighted; the
    /// router prefers whichever model produces a first token soonest.
    /// On macOS 26+ this typically resolves to `apple/foundation:default`
    /// via the Foundation Models system-LLM bonus. Reached via the
    /// `IntentHint::prefer_fast` flag (or `RoutingWorkload::Fastest`
    /// directly when callers know they want it).
    Fastest,
    /// Quality-critical, infrequent work where the BEST available model
    /// matters far more than latency or cost — e.g. building/verifying an
    /// agent, deriving a contract, structured extraction that a weak model
    /// gets wrong. The mirror image of `Fastest`: quality dominates, latency
    /// and cost are near-floor, so the most capable candidate wins decisively
    /// instead of the cheapest one squeaking past on cost. Reached via the
    /// `IntentHint::prefer_quality` flag.
    Quality,
}

impl RoutingWorkload {
    pub fn is_latency_sensitive(self) -> bool {
        matches!(
            self,
            RoutingWorkload::Interactive
                | RoutingWorkload::LocalPreferred
                | RoutingWorkload::Fastest,
        )
    }

    pub fn weights(self) -> (f64, f64, f64) {
        // Tuple is `(quality, latency, cost)` — destructured at the
        // single use site `adaptive_router.rs:892`.
        match self {
            RoutingWorkload::Interactive => (0.45, 0.40, 0.15),
            RoutingWorkload::Batch => (0.60, 0.15, 0.25),
            RoutingWorkload::Background => (0.65, 0.05, 0.30),
            // Quality-aware (closer to Interactive) but tolerant of
            // some latency hit since the caller chose local. Cost
            // weight matches Batch.
            RoutingWorkload::LocalPreferred => (0.55, 0.20, 0.25),
            // Voice fast track: latency is everything. Quality and
            // cost are deliberately near-floor — first audio in
            // <500ms beats any quality gain that takes another
            // round-trip. Sums to 1.0 like every other variant.
            RoutingWorkload::Fastest => (0.10, 0.85, 0.05),
            // Quality dominates; latency and cost near-floor so the most
            // capable model wins outright (the inverse of Fastest). Sums to 1.0.
            RoutingWorkload::Quality => (0.85, 0.05, 0.10),
        }
    }

    pub fn local_bonus(self) -> f64 {
        match self {
            RoutingWorkload::Interactive => 0.0,
            RoutingWorkload::Batch => 0.08,
            RoutingWorkload::Background => 0.15,
            // Stronger push than Background — "prefer local" should
            // win ties decisively, otherwise the hint is ineffective.
            RoutingWorkload::LocalPreferred => 0.20,
            // Local inference avoids network round-trips entirely —
            // the single biggest latency win available. On macOS the
            // Foundation Models `system_llm_bonus` stacks on top of
            // this for `apple/foundation:default`. Match
            // LocalPreferred's bonus so cloud-streamed fast paths
            // (e.g. gpt-4o-mini) can still win when locally there's
            // no model loaded; the weight profile already strongly
            // favours latency.
            RoutingWorkload::Fastest => 0.20,
            // Neutral: quality-first picks the best model wherever it lives.
            // A strong remote should win when available; on a local-only host
            // the biggest capable local wins on quality alone (no local bias
            // needed — or wanted, since it could mask a better cloud model).
            RoutingWorkload::Quality => 0.0,
        }
    }
}

/// Qwen3 hybrid thinking control. Qwen3 models were trained with both a
/// "thinking" (chain-of-thought inside `<think>...</think>`) and a
/// non-thinking mode. Upstream defaults thinking ON; `/no_think` and
/// `/think` are the documented per-turn overrides in the chat template.
///
/// Scope: applies to the *single-turn* local Qwen3 path driven by
/// [`apply_chat_template`]. The multi-turn `messages: Vec<Message>`
/// field on [`GenerateRequest`] is consumed by remote protocol
/// handlers (OpenAI/Anthropic/Google) which pass through user-supplied
/// system messages verbatim; this flag is not injected there. If you
/// need Qwen3 thinking control over a remote API, include `/think` or
/// `/no_think` explicitly in your own system message.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum ThinkingMode {
    /// Let the model decide. No explicit `/think` or `/no_think` directive
    /// is injected into the system prompt, so Qwen3's trained default
    /// (thinking ON) applies. `<think>...</think>` output is stripped
    /// from the returned text.
    #[default]
    Auto,
    /// Inject `/think` into the system prompt to explicitly request the
    /// thinking phase. Useful when callers want to force reasoning even
    /// on short prompts the model would normally answer directly.
    On,
    /// Inject `/no_think` into the system prompt to suppress the
    /// thinking phase for faster, more direct responses. This was the
    /// prior hard-coded behavior; callers now opt into it explicitly.
    Off,
}

impl ThinkingMode {
    /// Return the directive marker to append to the system prompt, or
    /// `None` when `Auto` (don't inject anything — trust model default).
    pub fn directive(self) -> Option<&'static str> {
        match self {
            ThinkingMode::Auto => None,
            ThinkingMode::On => Some("/think"),
            ThinkingMode::Off => Some("/no_think"),
        }
    }
}

/// Parameters controlling generation behavior.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerateParams {
    /// Sampling temperature (0.0 = greedy, 1.0 = full distribution).
    #[serde(default = "default_temperature")]
    pub temperature: f64,
    /// Hard-pin: when true, do NOT silently degrade a remote-only model chain to
    /// an installed on-device model as a last resort (see
    /// `generate_tracked_inner`). A caller that pinned a specific backbone — the
    /// coder's `--model`, or an A/B arm whose validity depends on the exact model
    /// — wants a hard error on a remote outage, not a silent swap to a weaker
    /// local model that manufactures fake results. Defaults false, so every other
    /// caller keeps the resilient degrade-to-local behavior. Lives on
    /// `GenerateParams` (not `GenerateRequest`) so the many explicit
    /// `GenerateRequest { .. }` initializers inherit it via `params: Default`.
    #[serde(default)]
    pub strict_model: bool,
    /// Top-p (nucleus) sampling threshold.
    #[serde(default = "default_top_p")]
    pub top_p: f64,
    /// Top-k sampling (0 = disabled).
    #[serde(default)]
    pub top_k: usize,
    /// Maximum tokens to generate.
    #[serde(default = "default_max_tokens")]
    pub max_tokens: usize,
    /// Stop sequences — generation halts when any is produced.
    #[serde(default)]
    pub stop: Vec<String>,
    /// Extended thinking budget (tokens). When > 0, enables the model's
    /// internal reasoning/planning phase before responding. Only supported
    /// by models with the ExtendedThinking capability (e.g., Claude).
    #[serde(default)]
    pub budget_tokens: usize,
    /// Routing workload class. Interactive requests bias toward lower latency,
    /// while batch/background work can tolerate slower high-quality local models.
    #[serde(default)]
    pub workload: RoutingWorkload,
    /// Tool choice mode, honored per-provider (see `protocol.rs`):
    /// * `"auto"` (default when tools present) / `"required"` / `"none"` — the
    ///   portable subset. Works on OpenAI-family, Anthropic, and Google
    ///   (`"required"` maps to Anthropic `any` / Google `ANY`).
    /// * `"any"` or a specific tool NAME — forcing modes honored on Anthropic
    ///   (`{"type":"any"}` / `{"type":"tool","name":…}`), Google (`ANY` +
    ///   `allowedFunctionNames`), and Bedrock (named only). **OpenAI-family
    ///   passes the string through verbatim**, so `"any"`/a bare name 400s
    ///   there — use `"required"` for portable forcing.
    /// * Caveats: Bedrock Converse has no `none` mode (a `"none"` string falls
    ///   through to a forced tool literally named "none"); `tool_choice` is only
    ///   emitted when `tools` are present. On Anthropic a forcing choice
    ///   (`any`/named) additionally disables extended thinking for that request.
    #[serde(default)]
    pub tool_choice: Option<String>,
    /// OpenAI-compatible parallel tool call control.
    #[serde(default)]
    pub parallel_tool_calls: Option<bool>,
    /// Qwen3 hybrid thinking mode control. `Auto` (default) leaves the
    /// model at its trained default (thinking on). `On`/`Off` inject
    /// the documented `/think` or `/no_think` directive into the chat
    /// template. Ignored by non-Qwen3 models.
    #[serde(default)]
    pub thinking: ThinkingMode,
    /// TTL for Anthropic prompt-cache entries, applied to every cache
    /// breakpoint this request emits (only when `cache_control` is enabled on
    /// the [`GenerateRequest`]). Defaults to the 5-minute ephemeral cache.
    /// Set [`CacheTtl::OneHour`] for agentic flows where minutes can pass
    /// between calls (tool execution, HITL approval gates) so the cached
    /// prefix survives instead of silently expiring into a full re-bill.
    #[serde(default)]
    pub cache_ttl: CacheTtl,
    /// Caller estimate of prompt/input tokens expected to be served from a
    /// provider cache. This is routing metadata only: it does not change the
    /// provider request or claim that a cache hit occurred. Defaults to zero
    /// because CAR cannot know a hit before the provider reports usage.
    #[serde(default)]
    pub estimated_cache_read_input_tokens: usize,
    /// Caller estimate of prompt/input tokens expected to be written into a
    /// provider cache. This is routing metadata only and defaults to zero;
    /// enabling `cache_control` alone never invents a cache-write estimate.
    #[serde(default)]
    pub estimated_cache_write_input_tokens: usize,
}

/// Prompt-cache time-to-live for Anthropic `cache_control` breakpoints.
///
/// The runtime applies one TTL uniformly to all breakpoints in a request, so
/// there is never a mix of 5-minute and 1-hour entries (Anthropic requires
/// 1-hour entries to precede 5-minute ones — a uniform TTL sidesteps the
/// ordering rule entirely).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheTtl {
    /// Default 5-minute ephemeral cache. Write billed at ~1.25× base input.
    #[default]
    FiveMinutes,
    /// Extended 1-hour cache. Write billed at ~2× base input. Best when the
    /// cached prefix is reused less often than every 5 minutes but more than
    /// hourly — typical of long agentic loops.
    OneHour,
}

impl CacheTtl {
    /// The value Anthropic expects in `cache_control.ttl`, or `None` for the
    /// default 5-minute cache (which omits the field entirely).
    pub fn anthropic_ttl(self) -> Option<&'static str> {
        match self {
            CacheTtl::FiveMinutes => None,
            CacheTtl::OneHour => Some("1h"),
        }
    }
}

fn default_temperature() -> f64 {
    0.7
}
fn default_top_p() -> f64 {
    0.9
}
/// The library default per-turn output budget. Used as the "unset" sentinel
/// by [`crate::InferenceEngine::generate_tracked`]: when a request still carries
/// this value, the engine substitutes the resolved model's real output ceiling
/// (`ModelSchema::effective_max_output()`) so long-horizon tool_use JSON isn't
/// truncated at 4096.
pub const DEFAULT_MAX_TOKENS: usize = 4096;
fn default_max_tokens() -> usize {
    DEFAULT_MAX_TOKENS
}

impl Default for GenerateParams {
    fn default() -> Self {
        Self {
            temperature: default_temperature(),
            strict_model: false,
            top_p: default_top_p(),
            top_k: 0,
            max_tokens: default_max_tokens(),
            stop: Vec::new(),
            budget_tokens: 0,
            workload: RoutingWorkload::Interactive,
            tool_choice: None,
            parallel_tool_calls: None,
            thinking: ThinkingMode::default(),
            cache_ttl: CacheTtl::default(),
            estimated_cache_read_input_tokens: 0,
            estimated_cache_write_input_tokens: 0,
        }
    }
}

// `Message`, `ToolCall`, and `ContentBlock` now live in the dependency-light
// `car-inference-types` crate so a consumer (car-sync's transcript-resume
// projection) can build the REAL types without the Candle/MLX stack — drift is
// a compile error there, not a runtime `from_value::<Message>` break. Re-exported
// here so every car-inference caller (and `crate::tasks::generate::Message`) is
// unchanged.
pub use car_inference_types::{ContentBlock, Message, Provenance, ThinkingBlock, ToolCall};

/// Constraint on the model's response shape. Distinct from `tools` —
/// tools are a side-channel for action invocation; `response_format`
/// constrains the *primary* text output to be parseable JSON, optionally
/// against a caller-supplied schema.
///
/// Provider mapping (handled in `protocol.rs`):
/// * **OpenAI / Azure / OpenAI-compatible**: `response_format: {type: "json_schema", json_schema: {schema, strict, name}}`
///   (or `{type: "json_object"}` for the looser variant). Strict mode rejects
///   any deviation from the schema.
/// * **Google (Gemini)**: `response_mime_type: "application/json"` plus
///   optional `response_schema`.
/// * **Anthropic**: not wired for a provider-enforced response-format field
///   under CAR's pinned `anthropic-version`. Both variants are rejected up
///   front with `UnsupportedMode` rather than silently weakening the caller's
///   output contract. Callers needing structured output on Claude should use
///   the `tools` + `tool_choice="required"` coercion idiom.
///
/// `JsonObject` is the looser variant — tells the provider "emit valid
/// JSON, no schema check required". Use when the schema is too dynamic
/// to spell out but the parse contract still matters.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ResponseFormat {
    /// JSON output validated against the provided schema. `strict: true`
    /// asks the provider to reject any deviation; `false` makes the
    /// schema a best-effort hint. `name` is OpenAI-specific (the
    /// `json_schema.name` field, max 64 chars, alphanumerics + `-_`);
    /// other providers ignore it.
    JsonSchema {
        schema: serde_json::Value,
        #[serde(default)]
        strict: bool,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        name: Option<String>,
    },
    /// Plain JSON-mode output — the provider emits valid JSON without
    /// schema enforcement.
    JsonObject,
}

/// A text generation request.
///
/// `Default` is derived so call sites can mutate just the fields
/// they care about: `GenerateRequest { prompt: "...".into(), ..Default::default() }`.
/// The default `prompt = ""` is useless on its own — callers always
/// override it — but the `..Default::default()` shorthand stops the
/// per-call-site mechanical churn every time a new optional field
/// lands (closes #109).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GenerateRequest {
    /// The prompt to complete (first user message for single-turn).
    pub prompt: String,
    /// Optional model override.
    pub model: Option<String>,
    /// Generation parameters.
    #[serde(default)]
    pub params: GenerateParams,
    /// Optional memory context to prepend to the prompt.
    /// When provided, this is injected as a system-level context block
    /// before the user prompt, grounding the model's response.
    #[serde(default)]
    pub context: Option<String>,
    /// Optional cache-breakpoint hint: the leading prefix of `context` that is
    /// stable across queries (Identity + Constraints, from the memory engine's
    /// `ContextSplit`). When set with `cache_control` on an Anthropic request,
    /// the system prompt is split into a cached stable block + an uncached
    /// volatile block so the stable prefix actually hits the cache. `None`
    /// keeps the whole-system single-block behavior. Must be a genuine prefix
    /// of `context` or it is ignored (the handler validates with `starts_with`).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub context_stable_prefix: Option<String>,
    /// Optional tool definitions for structured tool_use.
    /// When provided, the model may return tool_calls instead of text.
    /// Each tool is a JSON object with: name, description, parameters (JSON Schema).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<serde_json::Value>>,
    /// Optional images for vision models.
    /// When provided with a single-turn prompt, these are included as image content blocks
    /// in the user message. For multi-turn with `messages`, use `UserMultimodal` variants instead.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub images: Option<Vec<ContentBlock>>,
    /// Optional multi-turn conversation history.
    /// When provided, the backend builds a proper multi-turn message array
    /// instead of a single user message. The `prompt` field is ignored when
    /// messages are present.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub messages: Option<Vec<Message>>,
    /// Enable prompt caching for Anthropic API.
    /// When true, system prompt and tools are marked with cache_control breakpoints,
    /// enabling cache reuse across parent/child agent calls sharing the same prefix.
    #[serde(default)]
    pub cache_control: bool,
    /// Constrain output to JSON (optionally schema-validated). See
    /// [`ResponseFormat`] for the per-provider mapping. Defaults to
    /// `None` — free-form text.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
    /// Caller-supplied routing intent. None preserves the existing
    /// adaptive vs. pinned-model behavior. When `Some`, the adaptive
    /// router uses the hint to filter candidates (hard `require`),
    /// override task selection, and bias the score profile
    /// (`prefer_local`). See [`crate::intent::IntentHint`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub intent: Option<crate::intent::IntentHint>,
    /// Opaque caller-supplied correlation token, echoed VERBATIM in the
    /// `inference.runner.invoke` payload and otherwise ignored by CAR.
    ///
    /// Exists because a delegated-inference host with more than one call in
    /// flight had no way to map an invoke back to its own request state
    /// (Parslee-ai/car-releases#78). `call_id` is minted by the daemon *after*
    /// the host called `infer`, and `GenerateRequest` is a typed struct, so
    /// extra fields were dropped on the serde round-trip — leaving hosts to
    /// smuggle a UUID through `prompt`. That worked only because delegated
    /// models ignore `prompt`, and would break silently the moment one stopped
    /// ignoring it.
    ///
    /// CAR never reads, routes on, or logs this as anything but an opaque
    /// string: it is the host's own identifier, not a CAR one. Not a
    /// substitute for `call_id`, which remains the key for
    /// `inference.runner.event` / `.complete` / `.fail`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub client_ref: Option<String>,
    /// Require the resolved immutable catalog row to match this SHA-256
    /// digest. A mismatch fails before any provider, runner, or local backend
    /// is dispatched and is retained across routed retries and fallbacks.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expected_row_digest: Option<String>,
    /// Optionally require the route to use this exact catalog revision. This
    /// protects a caller that selected a row from a previously read snapshot
    /// from silently running against a different catalog generation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expected_catalog_revision: Option<String>,
    /// Who originated this request, set IN-PROCESS by the daemon and never
    /// carried on the wire (`#[serde(skip)]`, so it is absent from the
    /// `inference.runner.invoke` payload and from any persisted request).
    ///
    /// Exists so a delegated call can be dispatched to the runner registered by
    /// the session that made it, instead of to whichever runner registered most
    /// recently (Parslee-ai/car-releases#77). The runner slot used to be a
    /// process-global that the last registrant overwrote, so two hosts sharing
    /// one daemon stole each other's calls.
    ///
    /// Distinct from [`Self::client_ref`], which is the HOST's own opaque token
    /// and is echoed to the host; this is CAR's own routing key and is never
    /// echoed. A caller that leaves it `None` gets the global-fallback runner,
    /// which is the pre-existing behaviour.
    #[serde(skip)]
    pub caller: Option<String>,
}

/// Wrap a raw prompt in Qwen3 chat format if it's not already formatted.
///
/// Thinking behavior follows the caller-supplied [`ThinkingMode`]:
/// * `Auto` — no directive injected; Qwen3's trained default (thinking
///   on) applies.
/// * `On` — the documented `/think` directive is appended to the system
///   message on its own line, and the model is allowed to emit a full
///   `<think>...</think>` block before its answer.
/// * `Off` — `/no_think` is appended to the system message *and* an
///   empty `<think>\n\n</think>` block is pre-filled after the assistant
///   marker, matching upstream Qwen3's `enable_thinking=False` jinja
///   template. The pre-filled closed tags structurally prevent the
///   model from emitting a thinking block even if the directive is
///   contradicted later in the prompt.
///
/// When context is provided it is injected into the system message to
/// ground the model's response with memory. The directive always
/// appears *after* the context blob so user-supplied memory cannot
/// nudge the directive's parse position.
pub fn apply_chat_template(prompt: &str, context: Option<&str>, thinking: ThinkingMode) -> String {
    if prompt.contains("<|im_start|>") {
        return prompt.to_string();
    }
    // Directive goes on its own line at the end of the system message
    // (never concatenated onto prose) so Qwen3's chat template parser
    // sees `/think`/`/no_think` as a standalone token, not as part of
    // "assistant. /no_think".
    let directive_line = match thinking.directive() {
        Some(d) => format!("\n{d}"),
        None => String::new(),
    };
    // For Off, pre-fill a closed empty thinking block after the
    // assistant marker. This mirrors upstream Qwen3's jinja behavior
    // when `enable_thinking=False` and is the hard-switch (structural)
    // form of the mode, whereas `/no_think` alone is a soft directive.
    let thinking_prefill = match thinking {
        ThinkingMode::Off => "<think>\n\n</think>\n\n",
        _ => "",
    };
    match context {
        Some(ctx) => format!(
            "<|im_start|>system\nYou are a helpful assistant. Use the following context to inform your response.\n\n{ctx}{directive_line}<|im_end|>\n\
             <|im_start|>user\n{prompt}<|im_end|>\n\
             <|im_start|>assistant\n{thinking_prefill}"
        ),
        None => format!(
            "<|im_start|>system\nYou are a helpful assistant.{directive_line}<|im_end|>\n\
             <|im_start|>user\n{prompt}<|im_end|>\n\
             <|im_start|>assistant\n{thinking_prefill}"
        ),
    }
}

/// Render a Qwen3 chat-format prompt from a multi-turn `messages` array
/// and/or `tools`, for the in-process MLX/candle backends — which, unlike the
/// remote/OpenAI-compatible path, have no native chat-template or tool API.
///
/// Mirrors Qwen3's documented (Hermes-style) tool convention: tool signatures
/// go in a `<tools></tools>` block inside the system turn, and the model is
/// instructed to emit `<tool_call>{"name":…,"arguments":…}</tool_call>`, which
/// [`parse_tool_calls`] extracts back out of the completion.
///
/// When neither `messages` nor `tools` is present this delegates verbatim to
/// [`apply_chat_template`], so the common single-turn text path is byte-for-byte
/// unchanged. The returned string ends in `<|im_start|>assistant\n` (plus the
/// thinking prefill, if any) ready for completion.
pub fn render_chat_prompt(req: &GenerateRequest) -> String {
    let has_msgs = req.messages.as_ref().is_some_and(|m| !m.is_empty());
    let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty());
    if !has_msgs && !has_tools {
        // Identical behavior to the legacy path — zero risk for non-tool calls.
        return apply_chat_template(&req.prompt, req.context.as_deref(), req.params.thinking);
    }
    // Already pre-formatted (e.g. rendered once upstream)? Pass through.
    if req.messages.is_none() && req.prompt.contains("<|im_start|>") {
        return req.prompt.clone();
    }

    let thinking = req.params.thinking;
    let directive_line = match thinking.directive() {
        Some(d) => format!("\n{d}"),
        None => String::new(),
    };
    let thinking_prefill = match thinking {
        ThinkingMode::Off => "<think>\n\n</think>\n\n",
        _ => "",
    };

    // --- system turn (folds any System message(s) + context + tools) ---
    let mut system_text = String::new();
    let mut had_system = false;
    if let Some(msgs) = &req.messages {
        for m in msgs {
            if let Message::System { content } = m {
                if had_system {
                    system_text.push_str("\n\n");
                }
                system_text.push_str(content);
                had_system = true;
            }
        }
    }
    if !had_system {
        system_text.push_str("You are a helpful assistant.");
    }
    if let Some(ctx) = &req.context {
        system_text.push_str("\n\n");
        system_text.push_str(ctx);
    }

    let mut out = String::new();
    out.push_str("<|im_start|>system\n");
    out.push_str(&system_text);
    if let Some(tools) = &req.tools {
        if !tools.is_empty() {
            out.push_str(&render_tools_block(tools));
        }
    }
    out.push_str(&directive_line);
    out.push_str("<|im_end|>\n");

    // --- conversation turns ---
    match &req.messages {
        Some(msgs) => {
            for m in msgs {
                match m {
                    // System already folded into the system turn above.
                    Message::System { .. } => {}
                    Message::User { content } => {
                        out.push_str("<|im_start|>user\n");
                        out.push_str(content);
                        out.push_str("<|im_end|>\n");
                    }
                    Message::UserMultimodal { content } => {
                        // The in-process text tower can't consume image/video/
                        // audio blocks (those requests reject upstream); render
                        // the text blocks so a mixed message still threads.
                        let text = content
                            .iter()
                            .filter_map(|b| match b {
                                ContentBlock::Text { text } => Some(text.as_str()),
                                _ => None,
                            })
                            .collect::<Vec<_>>()
                            .join("\n");
                        out.push_str("<|im_start|>user\n");
                        out.push_str(&text);
                        out.push_str("<|im_end|>\n");
                    }
                    Message::Assistant {
                        content,
                        tool_calls,
                        .. // local Qwen template has no thinking-block concept
                    } => {
                        out.push_str("<|im_start|>assistant\n");
                        out.push_str(content);
                        for tc in tool_calls {
                            let args = serde_json::to_string(&tc.arguments)
                                .unwrap_or_else(|_| "{}".to_string());
                            out.push_str(&format!(
                                "\n<tool_call>\n{{\"name\": \"{}\", \"arguments\": {}}}\n</tool_call>",
                                tc.name, args
                            ));
                        }
                        out.push_str("<|im_end|>\n");
                    }
                    Message::ToolResult { content, .. } => {
                        // Qwen wraps tool results in a user turn with
                        // <tool_response> tags.
                        out.push_str("<|im_start|>user\n<tool_response>\n");
                        out.push_str(content);
                        out.push_str("\n</tool_response><|im_end|>\n");
                    }
                    // Provider-specific opaque items don't round-trip to a
                    // local model — drop them (matches the builder contract).
                    Message::ProviderOutputItems { .. } => {}
                }
            }
        }
        None => {
            // Single-turn with tools: the prompt is the user message.
            out.push_str("<|im_start|>user\n");
            out.push_str(&req.prompt);
            out.push_str("<|im_end|>\n");
        }
    }

    out.push_str("<|im_start|>assistant\n");
    out.push_str(thinking_prefill);
    out
}

/// Render the Qwen3 `<tools>` system-prompt block from tool definitions. Each
/// tool is normalized to `{"type":"function","function":{…}}` (Qwen's expected
/// shape); a definition already in that shape is passed through unchanged.
fn render_tools_block(tools: &[serde_json::Value]) -> String {
    let mut s = String::from(
        "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n\
         You are provided with function signatures within <tools></tools> XML tags:\n<tools>",
    );
    for t in tools {
        let is_wrapped = t.get("type").and_then(|v| v.as_str()) == Some("function")
            && t.get("function").is_some();
        let func = if is_wrapped {
            t.clone()
        } else {
            serde_json::json!({ "type": "function", "function": t })
        };
        s.push('\n');
        s.push_str(&serde_json::to_string(&func).unwrap_or_default());
    }
    s.push_str(
        "\n</tools>\n\nFor each function call, return a json object with function name and \
         arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n\
         {\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>",
    );
    s
}

/// Extract Qwen3 `<tool_call>{json}</tool_call>` blocks from a completion.
///
/// Returns the text with those blocks removed (trimmed) plus the parsed calls.
/// Each block is `{"name": …, "arguments": {…}}`; ids are minted locally
/// (`call_0`, `call_1`, …) since a local model emits none. A malformed or
/// unterminated block is left in the text and skipped, so a stray `<tool_call>`
/// literal never yields a bogus call. This is the local counterpart to the
/// structured `tool_calls` the remote `protocol.rs` path parses.
/// Strip a leaked model reasoning-channel prefix from answer text.
///
/// Table-stakes chat UIs (Claude Code, Codex, ChatGPT) never show the model's
/// internal reasoning label as the answer. Some models leak one: gemma-4 renders
/// reasoning as `<|channel>thought…<channel|>`, and with `skip_special_tokens`
/// the marker tokens drop but the plain channel name `thought` survives as a
/// leading token (observed `"thought\n<answer>"` and `"thought <answer>"`).
/// This runs on the finalized completion in BOTH the streaming (`parse_tool_calls`)
/// and non-streaming (backend `parse_tool_calls`) paths, so the chat surface,
/// the FFI harnesses, and memory all see clean text. Conservative: only a
/// *leading* bare `thought` marker followed by whitespace is removed, so a normal
/// answer that merely contains the word is untouched.
pub fn strip_leaked_reasoning(text: &str) -> String {
    // Delimited span, if the marker tokens survived decoding.
    if let Some(start) = text.find("<|channel>thought") {
        if let Some(rel) = text[start..].find("<channel|>") {
            let end = start + rel + "<channel|>".len();
            let joined = format!("{}{}", &text[..start], &text[end..]);
            return strip_leaked_reasoning(joined.trim());
        }
    }
    // Leading bare `thought` marker (delimiter tokens stripped by the decoder),
    // followed by any whitespace (newline OR space, both observed).
    let trimmed = text.trim_start();
    if let Some(rest) = trimmed.strip_prefix("thought") {
        if rest.starts_with(char::is_whitespace) {
            return rest.trim_start().to_string();
        }
    }
    text.to_string()
}

pub fn parse_tool_calls(text: &str) -> (String, Vec<ToolCall>) {
    const OPEN: &str = "<tool_call>";
    const CLOSE: &str = "</tool_call>";
    if !text.contains(OPEN) {
        return (strip_leaked_reasoning(text), Vec::new());
    }
    let mut calls = Vec::new();
    let mut clean = String::new();
    let mut rest = text;
    let mut idx = 0usize;
    while let Some(start) = rest.find(OPEN) {
        let after_open = start + OPEN.len();
        if let Some(end_rel) = rest[after_open..].find(CLOSE) {
            let json_end = after_open + end_rel;
            let block = rest[after_open..json_end].trim();
            if let Some(tc) = parse_one_tool_call(block, &mut idx) {
                clean.push_str(&rest[..start]);
                calls.push(tc);
                rest = &rest[json_end + CLOSE.len()..];
                continue;
            }
        }
        // Unterminated or unparseable — keep the opener literal, advance past it.
        clean.push_str(&rest[..after_open]);
        rest = &rest[after_open..];
    }
    clean.push_str(rest);
    (strip_leaked_reasoning(clean.trim()), calls)
}

fn parse_one_tool_call(block: &str, idx: &mut usize) -> Option<ToolCall> {
    let v: serde_json::Value = serde_json::from_str(block).ok()?;
    let name = v.get("name")?.as_str()?.to_string();
    let arguments = match v.get("arguments") {
        Some(serde_json::Value::Object(m)) => m.clone().into_iter().collect(),
        // Some models emit arguments as a JSON-encoded string.
        Some(serde_json::Value::String(s)) => serde_json::from_str(s).unwrap_or_default(),
        _ => std::collections::HashMap::new(),
    };
    let id = Some(format!("call_{}", *idx));
    *idx += 1;
    Some(ToolCall {
        id,
        name,
        arguments,
    })
}

/// Strip Qwen3 `<think>...</think>` blocks from model output, honoring
/// the caller's requested [`ThinkingMode`]:
///
/// * `On` — the caller explicitly asked for reasoning; return the raw
///   text verbatim so `<think>...</think>` is visible.
/// * `Auto` / `Off` — strip the thinking block and return only the
///   post-thinking answer. If the output contains an opening `<think>`
///   without a closing tag (truncation or stop before the model
///   finished thinking) return an empty string rather than leaking a
///   dangling tag to the caller.
/// Cut `text` at the earliest occurrence of any stop sequence (exclusive).
/// The generation loops detect a stop sequence by substring match and then
/// `break`, but the matched tokens are already in `generated`, so the decoded
/// text still contained the stop string — standard completion semantics
/// exclude it. Empty stop entries are ignored. No-op when nothing matches.
pub fn truncate_at_stop(text: &str, stops: &[String]) -> String {
    let mut cut = text.len();
    for s in stops {
        if s.is_empty() {
            continue;
        }
        if let Some(idx) = text.find(s.as_str()) {
            cut = cut.min(idx);
        }
    }
    text[..cut].to_string()
}

pub fn strip_thinking(text: &str, thinking: ThinkingMode) -> String {
    if matches!(thinking, ThinkingMode::On) {
        return text.to_string();
    }
    strip_thinking_block(text)
}

/// Remove a leading `<think>...</think>` block unconditionally.
/// Returns "" if `<think>` opens but never closes (incomplete output).
///
/// When that "opened but never closed" branch fires, log a warn line
/// — the caller is about to receive an empty string for what was
/// almost certainly a budget-truncation. Surfaces issue #168's root
/// cause without changing the return contract: callers (e.g. car-cli)
/// that look at stderr can tell users to either bump
/// `--max-tokens` or pass `--thinking off`. The decision lives in
/// the strip helper because every text-completion path funnels
/// through it; logging at the call sites would be a lot of
/// duplication.
fn strip_thinking_block(text: &str) -> String {
    if let Some(end) = text.find("</think>") {
        text[end + 8..].trim_start().to_string()
    } else if text.contains("<think>") {
        tracing::warn!(
            target: "car_inference::tasks::generate",
            raw_len = text.len(),
            "model output opened <think> but never closed it — \
             likely truncated by max_tokens; returning empty text. \
             Increase max_tokens, or set thinking=off to suppress \
             the reasoning phase."
        );
        String::new()
    } else {
        text.to_string()
    }
}

/// Callback for FLARE-style re-retrieval during generation.
/// Called with partial generation text, returns additional context or None.
pub type RetrievalCallback = Box<dyn Fn(&str) -> Option<String> + Send>;

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
/// Generate text from a prompt using the loaded model.
///
/// Returns `(text, time_to_first_token_ms, prompt_tokens, completion_tokens)`.
/// TTFT is measured from function entry through prefill to the moment the first
/// generated token has been sampled — the user-visible "did anything happen
/// yet" gate. `None` only when the prompt encodes to zero tokens (degenerate
/// input).
///
/// The token counts are reported so this path can populate `TokenUsage` like
/// the MLX and remote paths do. It previously returned only `(text, ttft)`, and
/// the caller hardcoded `usage: None` — so on every non-Apple-Silicon platform
/// (Linux, Windows, and macOS built with `car_skip_mlx`) EVERY local inference
/// reported no usage at all. A consumer summing `usage.total_tokens` read a
/// silent zero, which is worse than an error because it looks like a valid
/// answer (Parslee-ai/car#795). `prompt_tokens` is post-truncation — what the
/// model actually saw, not what the caller sent.
pub async fn generate(
    backend: &mut CandleBackend,
    req: GenerateRequest,
) -> Result<(String, Option<u64>, usize, usize), InferenceError> {
    let start = std::time::Instant::now();

    // Reset KV cache so each generation starts fresh (prevents cross-call state bleed)
    backend.clear_kv_cache();

    let formatted = apply_chat_template(&req.prompt, req.context.as_deref(), req.params.thinking);
    let tokens = backend.encode(&formatted)?;
    let eos = backend.eos_token_id();
    let eos_alt = backend.token_id("<|im_end|>");
    let params = &req.params;

    if tokens.is_empty() {
        return Ok((String::new(), None, 0, 0));
    }

    // Truncate to model's max context length minus generation headroom.
    // This prevents KV cache overflow on long prompts.
    let max_ctx = backend.context_length().unwrap_or(32768);
    let headroom = params.max_tokens.min(max_ctx / 4);
    let max_prompt = max_ctx.saturating_sub(headroom);
    let tokens = if tokens.len() > max_prompt {
        eprintln!(
            "[car-inference] truncating prompt from {} to {} tokens (context_length={})",
            tokens.len(),
            max_prompt,
            max_ctx
        );
        tokens[tokens.len() - max_prompt..].to_vec()
    } else {
        tokens
    };

    let mut generated = Vec::new();

    // Prefill: process all prompt tokens, sample first generated token from prefill logits
    let logits = backend.forward(&tokens, 0)?;
    let mut next_token = sample_token(&logits, params)?;
    let ttft_ms = Some(start.elapsed().as_millis() as u64);

    for _i in 0..params.max_tokens {
        // Check EOS
        if (eos == Some(next_token)) || (eos_alt == Some(next_token)) {
            break;
        }

        generated.push(next_token);

        // Check stop sequences
        if !params.stop.is_empty() {
            let text_so_far = backend.decode(&generated)?;
            if params.stop.iter().any(|s| text_so_far.contains(s)) {
                break;
            }
        }

        // Generate next token
        let pos = tokens.len() + generated.len() - 1;
        let logits = backend.forward(&[next_token], pos)?;
        next_token = sample_token(&logits, params)?;
    }

    let text = truncate_at_stop(&backend.decode(&generated)?, &params.stop);
    Ok((
        strip_thinking(&text, params.thinking),
        ttft_ms,
        tokens.len(),
        generated.len(),
    ))
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
/// Generate with FLARE-style confidence-triggered re-retrieval.
///
/// Monitors token logit confidence during generation. When a window of
/// low-confidence tokens is detected, pauses, re-queries memory with the
/// partial generation, and resumes with augmented context.
pub async fn generate_with_retrieval(
    backend: &mut CandleBackend,
    mut req: GenerateRequest,
    retrieval_cb: RetrievalCallback,
) -> Result<String, InferenceError> {
    // First pass: generate normally
    backend.clear_kv_cache();
    let formatted = apply_chat_template(&req.prompt, req.context.as_deref(), req.params.thinking);
    let tokens = backend.encode(&formatted)?;
    let eos = backend.eos_token_id();
    let eos_alt = backend.token_id("<|im_end|>");
    let params = req.params.clone();

    if tokens.is_empty() {
        return Ok(String::new());
    }

    let mut generated = Vec::new();
    let mut low_confidence_count = 0u32;
    let mut retrieval_attempts = 0u32;
    let max_retrievals = 2;
    let confidence_threshold = 0.4f32;
    let low_confidence_window = 3u32;

    let logits = backend.forward(&tokens, 0)?;
    let mut next_token = sample_token(&logits, &params)?;

    for _i in 0..params.max_tokens {
        if (eos == Some(next_token)) || (eos_alt == Some(next_token)) {
            break;
        }

        generated.push(next_token);

        // Generate next token and check confidence
        let pos = tokens.len() + generated.len() - 1;
        let logits = backend.forward(&[next_token], pos)?;

        // Check max logit probability for confidence
        let logits_f32: Vec<f32> = logits
            .squeeze(0)
            .unwrap_or(logits.clone())
            .to_dtype(candle_core::DType::F32)
            .map_err(|e| InferenceError::InferenceFailed(format!("dtype: {e}")))?
            .to_vec1()
            .unwrap_or_default();

        if !logits_f32.is_empty() {
            // Compute softmax max probability
            let max_logit = logits_f32.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let exp_sum: f32 = logits_f32.iter().map(|&v| (v - max_logit).exp()).sum();
            let max_prob = 1.0 / exp_sum; // probability of the top token

            if max_prob < confidence_threshold {
                low_confidence_count += 1;
            } else {
                low_confidence_count = 0;
            }

            // Trigger re-retrieval after sustained low confidence
            if low_confidence_count >= low_confidence_window && retrieval_attempts < max_retrievals
            {
                retrieval_attempts += 1;
                low_confidence_count = 0;

                // Use partial generation as re-retrieval query
                let partial = backend.decode(&generated)?;
                if let Some(new_context) = retrieval_cb(&partial) {
                    // Restart generation with augmented context
                    let combined_context = match req.context.take() {
                        Some(old) => format!("{}\n\n{}", old, new_context),
                        None => new_context,
                    };
                    req.context = Some(combined_context);

                    // Re-encode and restart
                    backend.clear_kv_cache();
                    let new_formatted = apply_chat_template(
                        &req.prompt,
                        req.context.as_deref(),
                        req.params.thinking,
                    );
                    let new_tokens = backend.encode(&new_formatted)?;
                    generated.clear();

                    let logits = backend.forward(&new_tokens, 0)?;
                    next_token = sample_token(&logits, &params)?;
                    continue;
                }
            }
        }

        next_token = sample_token(&logits, &params)?;
    }

    let text = backend.decode(&generated)?;
    Ok(strip_thinking(&text, params.thinking))
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
/// Sample a token, suppressing specific token IDs (set to -inf before sampling).
pub fn sample_token_suppress(
    logits: &Tensor,
    params: &GenerateParams,
    suppress: &[u32],
) -> Result<u32, InferenceError> {
    if suppress.is_empty() {
        return sample_token(logits, params);
    }
    // Clone logits and set suppressed tokens to -inf
    let mut logits_vec: Vec<f32> = logits
        .squeeze(0)
        .unwrap_or(logits.clone())
        .to_dtype(candle_core::DType::F32)
        .map_err(|e| InferenceError::InferenceFailed(format!("dtype: {e}")))?
        .to_vec1()
        .map_err(|e| InferenceError::InferenceFailed(format!("to_vec: {e}")))?;
    // Handle 2D logits (take last row)
    let dims = logits.dims();
    if dims.len() == 2 {
        let vocab = dims[dims.len() - 1];
        let start = logits_vec.len() - vocab;
        logits_vec = logits_vec[start..].to_vec();
    }
    for &id in suppress {
        if (id as usize) < logits_vec.len() {
            logits_vec[id as usize] = f32::NEG_INFINITY;
        }
    }
    let modified = Tensor::from_vec(
        logits_vec,
        logits.squeeze(0).unwrap_or(logits.clone()).shape(),
        logits.device(),
    )
    .map_err(|e| InferenceError::InferenceFailed(format!("from_vec: {e}")))?;
    sample_token(&modified, params)
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
/// Sample a token from logits using temperature + top-p + top-k.
pub fn sample_token(logits: &Tensor, params: &GenerateParams) -> Result<u32, InferenceError> {
    let logits = logits
        .squeeze(0)
        .map_err(|e| InferenceError::InferenceFailed(format!("squeeze: {e}")))?;
    let logits = logits
        .to_dtype(candle_core::DType::F32)
        .map_err(|e| InferenceError::InferenceFailed(format!("dtype: {e}")))?;

    // Get last position's logits
    let dim = logits.dims();
    let logits = if dim.len() == 2 {
        logits
            .get(dim[0] - 1)
            .map_err(|e| InferenceError::InferenceFailed(format!("get last: {e}")))?
    } else {
        logits
    };

    // Greedy decoding
    if params.temperature <= 0.0 {
        let token = logits
            .argmax(0)
            .map_err(|e| InferenceError::InferenceFailed(format!("argmax: {e}")))?
            .to_scalar::<u32>()
            .map_err(|e| InferenceError::InferenceFailed(format!("scalar: {e}")))?;
        return Ok(token);
    }

    // Temperature scaling
    let logits = (&logits / params.temperature)
        .map_err(|e| InferenceError::InferenceFailed(format!("temp scale: {e}")))?;

    let mut logits_vec: Vec<f32> = logits
        .to_vec1()
        .map_err(|e| InferenceError::InferenceFailed(format!("to_vec: {e}")))?;

    // Top-k filtering
    if params.top_k > 0 && params.top_k < logits_vec.len() {
        let mut indexed: Vec<(usize, f32)> = logits_vec.iter().copied().enumerate().collect();
        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        let threshold = indexed[params.top_k].1;
        for v in &mut logits_vec {
            if *v < threshold {
                *v = f32::NEG_INFINITY;
            }
        }
    }

    // Softmax
    let max_logit = logits_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let exp: Vec<f32> = logits_vec.iter().map(|&v| (v - max_logit).exp()).collect();
    let sum: f32 = exp.iter().sum();
    let mut probs: Vec<f32> = exp.iter().map(|&v| v / sum).collect();

    // Top-p (nucleus) filtering
    if params.top_p < 1.0 {
        let mut sorted_indices: Vec<usize> = (0..probs.len()).collect();
        sorted_indices.sort_by(|&a, &b| {
            probs[b]
                .partial_cmp(&probs[a])
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let mut cumsum = 0.0f32;
        let mut cutoff_idx = sorted_indices.len();
        for (i, &idx) in sorted_indices.iter().enumerate() {
            cumsum += probs[idx];
            if cumsum > params.top_p as f32 {
                cutoff_idx = i + 1;
                break;
            }
        }

        let keep: std::collections::HashSet<usize> =
            sorted_indices[..cutoff_idx].iter().copied().collect();
        for (i, p) in probs.iter_mut().enumerate() {
            if !keep.contains(&i) {
                *p = 0.0;
            }
        }

        // Renormalize
        let sum: f32 = probs.iter().sum();
        if sum > 0.0 {
            for p in &mut probs {
                *p /= sum;
            }
        }
    }

    // Categorical sample
    let r: f32 = rand_f32();
    let mut cumsum = 0.0f32;
    for (i, &p) in probs.iter().enumerate() {
        cumsum += p;
        if cumsum >= r {
            return Ok(i as u32);
        }
    }

    // Fallback: return highest prob token
    Ok(probs
        .iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(i, _)| i as u32)
        .unwrap_or(0))
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
/// Random float in [0, 1) using the rand crate.
fn rand_f32() -> f32 {
    rand::random::<f32>()
}

#[cfg(test)]
mod tool_format_tests {
    use super::*;

    fn req_with(
        prompt: &str,
        tools: Option<Vec<serde_json::Value>>,
        messages: Option<Vec<Message>>,
    ) -> GenerateRequest {
        GenerateRequest {
            prompt: prompt.to_string(),
            tools,
            messages,
            ..Default::default()
        }
    }

    #[test]
    fn catalog_identity_preconditions_round_trip_additively() {
        let request: GenerateRequest = serde_json::from_value(serde_json::json!({
            "prompt": "bound route",
            "expected_row_digest": "a".repeat(64),
            "expected_catalog_revision": "b".repeat(64),
        }))
        .unwrap();
        let encoded = serde_json::to_value(request).unwrap();

        assert_eq!(encoded["expected_row_digest"], "a".repeat(64));
        assert_eq!(encoded["expected_catalog_revision"], "b".repeat(64));
    }

    #[test]
    fn no_tools_no_messages_delegates_to_apply_chat_template() {
        let req = req_with("hello", None, None);
        assert_eq!(
            render_chat_prompt(&req),
            apply_chat_template("hello", None, req.params.thinking)
        );
    }

    #[test]
    fn single_turn_with_tools_emits_tools_block_and_user_turn() {
        let tools = vec![serde_json::json!({
            "name": "get_weather",
            "description": "Get weather",
            "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}
        })];
        let out = render_chat_prompt(&req_with("weather in Denver?", Some(tools), None));
        assert!(out.contains("<tools>"));
        assert!(out.contains("</tools>"));
        assert!(out.contains("get_weather"));
        // tools are normalized to the function-wrapper shape
        assert!(out.contains("\"type\":\"function\""));
        assert!(out.contains("<|im_start|>user\nweather in Denver?<|im_end|>"));
        assert!(out.trim_end().ends_with("<|im_start|>assistant"));
    }

    #[test]
    fn already_wrapped_tool_is_not_double_wrapped() {
        let tools = vec![serde_json::json!({
            "type": "function",
            "function": {"name": "f", "description": "d", "parameters": {}}
        })];
        let out = render_chat_prompt(&req_with("hi", Some(tools), None));
        assert_eq!(out.matches("\"type\":\"function\"").count(), 1);
    }

    #[test]
    fn multi_turn_renders_tool_result_and_assistant_tool_call() {
        let mut args = std::collections::HashMap::new();
        args.insert("city".to_string(), serde_json::json!("Denver"));
        let messages = vec![
            Message::System {
                content: "You are helpful.".into(),
            },
            Message::User {
                content: "weather?".into(),
            },
            Message::Assistant {
                content: String::new(),
                tool_calls: vec![ToolCall {
                    id: Some("call_0".into()),
                    name: "get_weather".into(),
                    arguments: args,
                }],
                thinking: vec![],
                model_id: None,
                local_last_resort: false,
            },
            Message::ToolResult {
                tool_use_id: "call_0".into(),
                content: "{\"temp\":72}".into(),
                provenance: Default::default(),
            },
        ];
        let out = render_chat_prompt(&req_with("", None, Some(messages)));
        assert!(out.contains("<|im_start|>system\nYou are helpful."));
        assert!(out.contains("<tool_call>"));
        assert!(out.contains("get_weather"));
        assert!(out.contains("<tool_response>\n{\"temp\":72}\n</tool_response>"));
        assert!(out.trim_end().ends_with("<|im_start|>assistant"));
    }

    #[test]
    fn parse_extracts_single_tool_call_and_strips_tags() {
        let text = "Let me check.\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Denver\"}}\n</tool_call>";
        let (clean, calls) = parse_tool_calls(text);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].name, "get_weather");
        assert_eq!(
            calls[0].arguments.get("city").unwrap(),
            &serde_json::json!("Denver")
        );
        assert_eq!(calls[0].id.as_deref(), Some("call_0"));
        assert!(!clean.contains("<tool_call>"));
        assert_eq!(clean, "Let me check.");
    }

    #[test]
    fn parse_extracts_multiple_tool_calls_with_distinct_ids() {
        let text = "<tool_call>\n{\"name\": \"a\", \"arguments\": {}}\n</tool_call><tool_call>\n{\"name\": \"b\", \"arguments\": {}}\n</tool_call>";
        let (_clean, calls) = parse_tool_calls(text);
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].name, "a");
        assert_eq!(calls[1].name, "b");
        assert_eq!(calls[0].id.as_deref(), Some("call_0"));
        assert_eq!(calls[1].id.as_deref(), Some("call_1"));
    }

    #[test]
    fn parse_plain_text_yields_no_calls() {
        let (clean, calls) = parse_tool_calls("just a normal answer");
        assert!(calls.is_empty());
        assert_eq!(clean, "just a normal answer");
    }

    #[test]
    fn parse_malformed_block_is_left_intact_and_skipped() {
        let text = "before <tool_call>\nnot json\n</tool_call> after";
        let (clean, calls) = parse_tool_calls(text);
        assert!(calls.is_empty());
        // The unparseable block is not silently dropped.
        assert!(clean.contains("<tool_call>"));
    }

    #[test]
    fn parse_arguments_as_encoded_string() {
        let text =
            "<tool_call>\n{\"name\": \"f\", \"arguments\": \"{\\\"k\\\": 1}\"}\n</tool_call>";
        let (_clean, calls) = parse_tool_calls(text);
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].arguments.get("k").unwrap(), &serde_json::json!(1));
    }
}

#[cfg(test)]
mod thinking_tests {
    use super::*;

    #[test]
    fn auto_injects_no_directive_and_no_prefill() {
        let out = apply_chat_template("hi", None, ThinkingMode::Auto);
        assert!(!out.contains("/no_think"));
        assert!(!out.contains("/think"));
        assert!(!out.contains("<think>"));
        assert!(out.contains("<|im_start|>user\nhi<|im_end|>"));
    }

    /// The build_context → prompt seam: an assembled context string must reach
    /// the rendered prompt (in the system slot, ahead of the user question)
    /// alongside the "Use the following context" preamble.
    #[test]
    fn assembled_context_lands_in_system_slot_before_user_question() {
        let out = apply_chat_template(
            "user question",
            Some("ASSEMBLED-CONTEXT-MARKER"),
            ThinkingMode::Auto,
        );
        assert!(
            out.contains("Use the following context"),
            "context preamble must be present: {out}"
        );
        assert!(
            out.contains("ASSEMBLED-CONTEXT-MARKER"),
            "assembled context must reach the prompt: {out}"
        );
        let marker_at = out.find("ASSEMBLED-CONTEXT-MARKER").unwrap();
        let question_at = out.find("user question").unwrap();
        assert!(
            marker_at < question_at,
            "context must precede the user question: {out}"
        );
    }

    #[test]
    fn off_injects_no_think_on_own_line_and_prefills_empty_think() {
        let out = apply_chat_template("hi", None, ThinkingMode::Off);
        // Directive on its own line, not concatenated onto prose.
        assert!(out.contains("\n/no_think<|im_end|>"));
        assert!(!out.contains(" /no_think"));
        // Closed empty thinking block pre-filled after assistant marker
        // — the upstream jinja hard-switch for enable_thinking=False.
        assert!(out.contains("<|im_start|>assistant\n<think>\n\n</think>\n\n"));
    }

    #[test]
    fn on_injects_think_and_no_prefill() {
        let out = apply_chat_template("hi", None, ThinkingMode::On);
        assert!(out.contains("\n/think<|im_end|>"));
        assert!(!out.contains("/no_think"));
        assert!(!out.contains("<think>"));
    }

    #[test]
    fn pre_formatted_prompt_is_untouched() {
        let pre = "<|im_start|>system\ncustom<|im_end|>\n<|im_start|>user\nhi<|im_end|>";
        let out = apply_chat_template(pre, None, ThinkingMode::Off);
        assert_eq!(out, pre);
    }

    #[test]
    fn directive_appears_after_context_not_before() {
        let out = apply_chat_template("q?", Some("some memory"), ThinkingMode::Off);
        let ctx_idx = out.find("some memory").unwrap();
        let directive_idx = out.find("/no_think").unwrap();
        assert!(
            directive_idx > ctx_idx,
            "directive must appear after context so user memory cannot nudge the parse"
        );
    }

    #[test]
    fn default_params_is_auto() {
        assert_eq!(GenerateParams::default().thinking, ThinkingMode::Auto);
    }

    #[test]
    fn thinking_mode_serde_snake_case() {
        let json = serde_json::to_string(&ThinkingMode::Off).unwrap();
        assert_eq!(json, "\"off\"");
        let parsed: ThinkingMode = serde_json::from_str("\"on\"").unwrap();
        assert_eq!(parsed, ThinkingMode::On);
    }

    #[test]
    fn strip_preserves_thinking_when_on() {
        let text = "<think>reasoning here</think>the answer";
        let out = strip_thinking(text, ThinkingMode::On);
        assert_eq!(
            out, text,
            "On mode must return raw text with <think> visible"
        );
    }

    #[test]
    fn strip_removes_thinking_when_auto_or_off() {
        let text = "<think>reasoning</think>the answer";
        assert_eq!(strip_thinking(text, ThinkingMode::Auto), "the answer");
        assert_eq!(strip_thinking(text, ThinkingMode::Off), "the answer");
    }

    #[test]
    fn strip_returns_empty_on_unterminated_think() {
        // Output was cut off mid-thinking — don't leak the dangling tag.
        let text = "<think>mid-reasoning, never closed";
        assert_eq!(strip_thinking(text, ThinkingMode::Auto), "");
        assert_eq!(strip_thinking(text, ThinkingMode::Off), "");
        // On mode still returns the raw text — caller asked for it.
        assert_eq!(strip_thinking(text, ThinkingMode::On), text);
    }

    #[test]
    fn strip_is_noop_when_no_think_tag() {
        let text = "just a plain answer";
        assert_eq!(strip_thinking(text, ThinkingMode::Auto), text);
        assert_eq!(strip_thinking(text, ThinkingMode::Off), text);
        assert_eq!(strip_thinking(text, ThinkingMode::On), text);
    }
}

#[cfg(test)]
mod workload_tests {
    use super::*;

    #[test]
    fn all_workload_weights_sum_to_one() {
        for w in [
            RoutingWorkload::Interactive,
            RoutingWorkload::Batch,
            RoutingWorkload::Background,
            RoutingWorkload::LocalPreferred,
            RoutingWorkload::Fastest,
        ] {
            let (q, l, c) = w.weights();
            let sum = q + l + c;
            assert!(
                (sum - 1.0).abs() < 1e-6,
                "weights for {w:?} sum to {sum}, expected 1.0"
            );
        }
    }

    #[test]
    fn fastest_weights_dominate_on_latency() {
        let (q, l, c) = RoutingWorkload::Fastest.weights();
        // Latency should be the largest by a wide margin — that's the
        // whole point of this workload class.
        assert!(l > q && l > c);
        assert!(l >= 0.7, "latency weight too small: {l}");
    }

    #[test]
    fn fastest_is_latency_sensitive() {
        assert!(RoutingWorkload::Fastest.is_latency_sensitive());
    }
}