unsloth/studio/backend/core/inference/llama_cpp.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""llama-server inference backend for GGUF models.

Manages a llama-server subprocess and proxies chat completions through its
OpenAI-compatible /v1/chat/completions endpoint.
"""

import atexit
import contextlib
import json
import os
import re
import struct
import structlog
from loggers import get_logger
import shutil
import signal
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Generator, Iterable, List, Optional

import httpx

from core.inference.llama_server_args import (
    parse_cache_override,
    parse_ctx_override,
    resolve_cache_type_kv,
    resolve_requested_ctx,
)
from core.tool_healing import (
    _TC_END_TAG_RE,
    _TC_FUNC_CLOSE_RE,
    _TC_FUNC_START_RE,
    _TC_JSON_START_RE,
    _TC_PARAM_CLOSE_RE,
    _TC_PARAM_START_RE,
    _TOOL_ALL_PATS,
    _TOOL_CLOSED_PATS,
    parse_tool_calls_from_text,
    strip_tool_call_markup,
)
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import (
    windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
from core.inference.tool_call_parser import (
    RAG_MAX_SEARCHES_PER_TURN,
    RAG_SEARCH_CAP_NUDGE,
    TOOL_XML_SIGNALS,
    parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
)
from core.inference.tool_loop_controller import (
    ToolLoopController,
    tool_event_provenance,
)

logger = get_logger(__name__)


# ── Pre-compiled patterns for plan-without-action re-prompt ──
# Forward-looking intent signals: the model is describing what it *will*
# do rather than giving a final answer.
_INTENT_SIGNAL = re.compile(
    r"(?i)("
    # Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes).
    # Excludes "I can"/"I should"/"I want to"/"let's" (common in answers).
    # Negative lookahead drops negated forms ("I will not") so a refusal
    # doesn't trigger a re-prompt.
    r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
    r"|"
    # Step/plan framing: "First ...", "Step 1:", "Here's my plan"
    r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
    r"|"
    # "Now I" / "Next I" patterns
    r"\b(?:now i|next i)\b"
    r")"
)
_MAX_REPROMPTS = 1

# Without max_tokens, llama-server defaults n_predict = n_ctx (up to 262144 for
# Qwen3.5), causing many-minute zombie decodes when cancel fails.
# t_max_predict_ms is a wall-clock backstop but per the llama.cpp README only
# fires after a newline, so we keep a token cap as the front-line limiter.
# The cap is the effective context length when known, else this floor. 4096 was
# too low: Qwen3 / gpt-oss reasoning traces and max_tokens-omitting OpenAI-API
# callers (langchain, llama-index, curl) got truncated mid-sentence.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_T_MAX_PREDICT_MS = 600_000  # 10 min
_REPROMPT_MAX_CHARS = 2000
_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
    r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
    re.I,
)
_FINAL_ANSWER_SIGNAL = re.compile(
    r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:)\b",
    re.I,
)


def _is_short_intent_without_action(text: str) -> bool:
    stripped = text.strip()
    return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None


def _should_suppress_forced_no_tool_output(text: str) -> bool:
    """Suppress only repeated forced-turn planning text, not final answers."""
    stripped = text.strip()
    if not stripped or len(stripped) >= _REPROMPT_MAX_CHARS:
        return False
    if _FINAL_ANSWER_SIGNAL.search(stripped):
        return False
    return _FORCED_REPEAT_PLAN_SIGNAL.search(stripped) is not None


# ── Pre-compiled patterns for GGUF shard detection ───────────
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")


# ── Sliding-window-pattern resolver ───────────────────────────
# Resolves the per-layer SWA mask when a GGUF reports a sliding window but
# no `sliding_window_pattern` field. Tier order in `_resolve_swa_pattern`:
# GGUF metadata, on-disk cache, bootstrap dict below, transformers
# introspection, HF Hub config.json, legacy 1/4 fallback. Period N means
# layer i is SWA iff `(i + 1) % N != 0`, matching transformers. Skipped on
# purpose: phi3 (no key/val length in GGUF, window >= ctx anyway), qwen2
# family (converter strips sliding_window when use_sliding_window=False),
# mistral v0.1/v0.2 (all-SWA can't be a period).
_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = {
    "gemma2": 2,  # Gemma2Config.sliding_window_pattern
    "gemma3": 6,  # Gemma3TextConfig.sliding_window_pattern
    "gemma3n": 5,  # text_config.layer_types: SWA*4 + FULL
    "gpt_oss": 2,  # text_config.layer_types: alternating
    "cohere2": 4,  # Cohere2Config.sliding_window_pattern
}

# Process-wide cache backed by JSON on disk. Values are int period or
# list[bool] mask. Lazy-loaded.
_SWA_CACHE: Optional[dict] = None
_SWA_CACHE_LOCK = threading.Lock()


def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool:
    """Quick DNS check on a daemon thread, so concurrent sockets aren't
    affected by socket.setdefaulttimeout."""
    result: list[Optional[bool]] = [None]

    def _probe() -> None:
        try:
            socket.gethostbyname(host)
            result[0] = False
        except Exception:
            result[0] = True

    t = threading.Thread(target = _probe, daemon = True)
    t.start()
    t.join(timeout)
    # Thread still running -> resolver wedged -> dead.
    return True if result[0] is None else result[0]


@contextlib.contextmanager
def _hf_offline_if_dns_dead():
    """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails;
    restores env on exit so a transient hiccup can't quarantine the process.
    No-op if the user already set it."""
    if "HF_HUB_OFFLINE" in os.environ:
        yield False
        return
    if not _probe_dns_dead():
        yield False
        return

    transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ
    os.environ["HF_HUB_OFFLINE"] = "1"
    if not transformers_was_set:
        os.environ["TRANSFORMERS_OFFLINE"] = "1"
    logger.warning("huggingface.co unreachable; using local HF cache for this load.")
    try:
        yield True
    finally:
        os.environ.pop("HF_HUB_OFFLINE", None)
        if not transformers_was_set:
            os.environ.pop("TRANSFORMERS_OFFLINE", None)


def _swa_cache_path() -> Path:
    home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
    base = Path(home) if home else Path.home() / ".unsloth" / "studio"
    return base / "swa_cache.json"


def _load_swa_cache() -> dict:
    global _SWA_CACHE
    with _SWA_CACHE_LOCK:
        if _SWA_CACHE is not None:
            return _SWA_CACHE
        try:
            with open(_swa_cache_path()) as f:
                _SWA_CACHE = json.load(f)
                if not isinstance(_SWA_CACHE, dict):
                    _SWA_CACHE = {}
        except (FileNotFoundError, json.JSONDecodeError, OSError):
            _SWA_CACHE = {}
        return _SWA_CACHE


def _save_swa_cache(cache: dict) -> None:
    try:
        path = _swa_cache_path()
        path.parent.mkdir(parents = True, exist_ok = True)
        tmp = path.with_suffix(".json.tmp")
        with open(tmp, "w") as f:
            json.dump(cache, f, indent = 2, sort_keys = True)
        tmp.replace(path)
    except OSError:
        pass


def _period_from_layer_types(layer_types: list) -> Optional[int]:
    """Smallest period N where `(i+1) % N != 0` matches the SWA mask, else None."""
    if not layer_types:
        return None
    is_swa = ["full" not in str(t).lower() for t in layer_types]
    n = len(is_swa)
    for N in range(1, n + 1):
        if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)):
            return N
    return None


def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
    try:
        from huggingface_hub import hf_hub_download
        cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
        with open(cfg_path) as f:
            cfg = json.load(f)
    except Exception:
        return None

    src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg
    period = src.get("sliding_window_pattern")
    if isinstance(period, int) and period > 0:
        return period
    lt = src.get("layer_types")
    if isinstance(lt, list) and lt:
        return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt]
    return None


def _arch_aliases(arch: str) -> tuple:
    # GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways.
    seen = []
    for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")):
        if a and a not in seen:
            seen.append(a)
    return tuple(seen)


def _swa_entry_from_config_obj(cfg) -> Optional[object]:
    src = getattr(cfg, "text_config", None) or cfg
    period = getattr(src, "sliding_window_pattern", None)
    if isinstance(period, int) and period > 0:
        return period
    lt = getattr(src, "layer_types", None)
    if isinstance(lt, list) and lt:
        return _period_from_layer_types(lt) or ["full" not in str(t).lower() for t in lt]
    return None


_SWA_PATTERN_SOURCE_RE = re.compile(r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)")


def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]:
    """Default-instantiate the matching Config; on failure, regex-parse its
    source for `sliding_window_pattern = N`."""
    try:
        from transformers.models.auto.configuration_auto import (
            CONFIG_MAPPING,
            CONFIG_MAPPING_NAMES,
        )
    except Exception:
        return None

    cfg_class = None
    for alias in _arch_aliases(arch):
        if alias in CONFIG_MAPPING_NAMES:
            try:
                cfg_class = CONFIG_MAPPING[alias]
                break
            except Exception:
                cfg_class = None
    if cfg_class is None:
        return None

    try:
        if (entry := _swa_entry_from_config_obj(cfg_class())) is not None:
            return entry
    except Exception:
        pass

    import inspect

    candidates = [cfg_class]
    text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config")
    if text_cfg_class is not None:
        candidates.append(text_cfg_class)
    for cls in candidates:
        try:
            src = inspect.getsource(cls)
        except (OSError, TypeError):
            continue
        if m := _SWA_PATTERN_SOURCE_RE.search(src):
            period = int(m.group(1))
            if period > 0:
                return period
    return None


def _resolve_swa_pattern(
    arch: Optional[str],
    n_layers: Optional[int],
    source_repo_candidates: tuple = (),
    *,
    allow_network: Optional[bool] = None,
) -> Optional[list]:
    if not arch or not n_layers:
        return None
    if allow_network is None:
        allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in (
            "1",
            "true",
            "True",
            "yes",
        )

    cache = _load_swa_cache()

    def _entry_to_mask(entry):
        if isinstance(entry, int) and entry > 0:
            return [(i + 1) % entry != 0 for i in range(n_layers)]
        if isinstance(entry, list) and entry:
            return [bool(entry[i % len(entry)]) for i in range(n_layers)]
        return None

    def _persist(entry):
        with _SWA_CACHE_LOCK:
            cache[arch] = entry
        _save_swa_cache(cache)

    if (entry := cache.get(arch)) is not None:
        if (mask := _entry_to_mask(entry)) is not None:
            return mask

    if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None:
        return _entry_to_mask(entry)

    entry = _resolve_swa_entry_from_transformers(arch)
    if entry is not None:
        _persist(entry)
        return _entry_to_mask(entry)

    # Tier 3: live HF fetch (result persistently cached)
    if allow_network:
        for repo_id in source_repo_candidates:
            if not repo_id:
                continue
            entry = _fetch_swa_entry_from_hf(repo_id)
            if entry is not None:
                _persist(entry)
                return _entry_to_mask(entry)

    return None


def _hf_repo_from_url(url: Optional[str]) -> Optional[str]:
    """Strip `https://huggingface.co/owner/name(/...)` -> `owner/name`."""
    if not url or "huggingface.co/" not in url:
        return None
    tail = url.split("huggingface.co/", 1)[1].rstrip("/")
    parts = tail.split("/")
    if len(parts) < 2:
        return None
    return f"{parts[0]}/{parts[1]}"


# Lazy import to avoid pulling transformers in at module level.
def _extract_model_size_b(model_id: str):
    from utils.models import extract_model_size_b
    return extract_model_size_b(model_id)


_TOOL_TEMPLATE_MARKERS = (
    "{%- if tools %}",
    "{%- if tools -%}",
    "{% if tools %}",
    "{% if tools -%}",
    '"role" == "tool"',
    "'role' == 'tool'",
    'message.role == "tool"',
    "message.role == 'tool'",
)


def detect_reasoning_flags(
    chat_template: Optional[str],
    model_identifier: Optional[str] = None,
    *,
    log_source: Optional[str] = None,
) -> dict:
    """Classify a chat template's reasoning and tool-calling capabilities.

    Returns the same five keys as the GGUF sniffer: ``supports_reasoning``,
    ``reasoning_style`` (``"enable_thinking"`` | ``"reasoning_effort"``),
    ``reasoning_always_on``, ``supports_preserve_thinking``,
    ``supports_tools``. Used by both the llama-server backend at load time
    and the safetensors/transformers paths in ``routes/inference`` so they
    agree on what the frontend sees.
    """
    flags = {
        "supports_reasoning": False,
        "reasoning_style": "enable_thinking",
        "reasoning_always_on": False,
        "supports_preserve_thinking": False,
        "supports_tools": False,
    }
    if not chat_template:
        return flags
    tpl = chat_template
    prefix = f"{log_source}: " if log_source else ""

    if "enable_thinking" in tpl:
        flags["supports_reasoning"] = True
        flags["reasoning_style"] = "enable_thinking"
        logger.info(f"{prefix}model supports reasoning (enable_thinking)")
    elif "reasoning_effort" in tpl:
        # gpt-oss / Harmony use reasoning_effort
        # ("low" | "medium" | "high"), not a boolean.
        flags["supports_reasoning"] = True
        flags["reasoning_style"] = "reasoning_effort"
        logger.info(f"{prefix}model supports reasoning (reasoning_effort)")
    elif "thinking" in tpl:
        # DeepSeek uses 'thinking', not 'enable_thinking'
        normalized_id = (model_identifier or "").lower()
        if "deepseek" in normalized_id:
            flags["supports_reasoning"] = True
            logger.info(f"{prefix}model supports reasoning (DeepSeek thinking)")

    # Hardcoded <think> tags or reasoning_content in the template mean
    # thinking is always on (no toggle).
    if not flags["supports_reasoning"]:
        if ("<think>" in tpl and "</think>" in tpl) or "reasoning_content" in tpl:
            flags["supports_reasoning"] = True
            flags["reasoning_always_on"] = True
            logger.info(f"{prefix}model always reasons (<think> tags in template)")

    # preserve_thinking: independent kwarg on some Qwen templates that
    # keeps historical <think> blocks in prior assistant turns.
    if "preserve_thinking" in tpl:
        flags["supports_preserve_thinking"] = True
        logger.info(f"{prefix}model supports preserve_thinking")

    if any(marker in tpl for marker in _TOOL_TEMPLATE_MARKERS):
        flags["supports_tools"] = True
        logger.info(f"{prefix}model supports tool calling")

    return flags


def _is_mtp_model_name(model_identifier: Optional[str], gguf_path: Optional[str] = None) -> bool:
    """Name-based MTP detector. Fallback for the metadata signal."""
    for cand in (model_identifier, Path(gguf_path).name if gguf_path else None):
        if cand and "-mtp" in cand.lower():
            return True
    return False


def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
    """User passed --spec-type / --spec-default? llama-server takes one
    --spec-type (comma-separated to chain), so suppress auto-emit."""
    if not extra_args:
        return False
    for raw in extra_args:
        tok = str(raw)
        if not tok.startswith("--"):
            continue
        flag = tok.split("=", 1)[0]
        if flag in ("--spec-type", "--spec-default"):
            return True
    return False


def _build_ngram_mod_flags(
    caps: Optional[dict],
    n_match: int = 24,
    n_min: int = 48,
    n_max: int = 64,
) -> list[str]:
    """Emit the right ngram-mod knob flags for the running llama-server.

    Post-rename builds expose ``--spec-ngram-mod-n-{match,min,max}``;
    pre-rename builds expose legacy ``--spec-ngram-size-n`` /
    ``--draft-min`` / ``--draft-max``. ``caps`` comes from
    ``probe_server_capabilities``; ``ngram_mod_flavor`` says which set is
    real (vs a removal-stub). Returns ``[]`` when neither is available so
    the caller can drop ngram-mod entirely.
    """
    flavor = caps.get("ngram_mod_flavor") if caps else None
    if flavor == "new":
        return [
            "--spec-ngram-mod-n-match",
            str(n_match),
            "--spec-ngram-mod-n-min",
            str(n_min),
            "--spec-ngram-mod-n-max",
            str(n_max),
        ]
    if flavor == "legacy":
        # Pre-rename llama.cpp: same knobs lived under --spec-ngram-size-n
        # (lookup length) and generic --draft-min / --draft-max (N range).
        return [
            "--spec-ngram-size-n",
            str(n_match),
            "--draft-min",
            str(n_min),
            "--draft-max",
            str(n_max),
        ]
    return []


# Canonical Speculative Decoding modes exposed by the Studio chat UI.
# Dropdown renders five (auto, mtp, ngram, mtp+ngram, off); the load API
# also accepts legacy values the original Switch and external callers emit
# (default, draft-mtp, ngram-mod, ngram-simple).
_CANONICAL_SPEC_MODES = {"auto", "mtp", "ngram", "mtp+ngram", "off", "ngram-simple"}
_LEGACY_SPEC_MODE_MAP = {
    "default": "auto",
    "draft-mtp": "mtp",
    "ngram-mod": "ngram",
}


def _canonicalize_spec_mode(value):
    """Map any accepted ``speculative_type`` input onto a canonical mode.

    Returns ``auto``, ``mtp``, ``ngram``, ``mtp+ngram``, ``off``,
    ``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``).
    Unknown strings collapse to ``auto`` so a stale UI value or typo falls
    back to the safe platform-aware path.
    """
    if value is None:
        return None
    if not isinstance(value, str):
        return None
    stripped = value.strip().lower()
    if not stripped:
        return None
    if stripped in _CANONICAL_SPEC_MODES:
        return stripped
    if stripped in _LEGACY_SPEC_MODE_MAP:
        return _LEGACY_SPEC_MODE_MAP[stripped]
    # Old persisted state emits llama.cpp comma-chains e.g.
    # "ngram-mod,draft-mtp"; collapse the most common one explicitly.
    pieces = [p.strip() for p in stripped.split(",") if p.strip()]
    has_mtp = any(p in ("mtp", "draft-mtp") for p in pieces)
    has_ngram = any(p in ("ngram", "ngram-mod") for p in pieces)
    if has_mtp and has_ngram:
        return "mtp+ngram"
    if has_mtp:
        return "mtp"
    if has_ngram:
        return "ngram"
    return "auto"


def _backfill_usage_from_timings(usage, timings):
    """Synthesize ``usage`` from llama-server's ``timings`` when the
    OpenAI-style usage block is missing or reports zero tokens.

    The Studio chat UI computes generation t/s from
    ``meta.usage.completion_tokens / totalStreamTime``. llama-server always
    populates ``timings.predicted_n`` (true decoded count) and
    ``timings.prompt_n``, but the final SSE chunk's ``usage`` can be absent
    or zero on some server builds / streaming configs, making the UI fall
    back to wall-clock t/s and dilute speculative-decoding speedups.
    """
    if not timings:
        return usage
    if usage and usage.get("completion_tokens"):
        return usage
    predicted_n = timings.get("predicted_n")
    prompt_n = timings.get("prompt_n")
    if predicted_n is None and prompt_n is None:
        return usage
    out = dict(usage or {})
    if not out.get("completion_tokens") and predicted_n is not None:
        out["completion_tokens"] = predicted_n
    if not out.get("prompt_tokens") and prompt_n is not None:
        out["prompt_tokens"] = prompt_n
    out["total_tokens"] = int(out.get("prompt_tokens") or 0) + int(
        out.get("completion_tokens") or 0
    )
    return out


class LlamaCppBackend:
    """Manages a llama-server subprocess for GGUF model inference.

    Lifecycle:
        1. load_model() — start llama-server with the GGUF file
        2. generate_chat_completion() — proxy to /v1/chat/completions, stream back
        3. unload_model() — terminate the subprocess
    """

    def __init__(self):
        self._process: Optional[subprocess.Popen] = None
        self._port: Optional[int] = None
        self._model_identifier: Optional[str] = None
        self._gguf_path: Optional[str] = None
        self._hf_repo: Optional[str] = None
        self._hf_variant: Optional[str] = None
        self._is_vision: bool = False
        self._healthy = False
        # Set by _classify_gpu_offload after _wait_for_health.
        self._gpu_offload_active: Optional[bool] = None
        self._context_length: Optional[int] = None
        self._effective_context_length: Optional[int] = None
        self._max_context_length: Optional[int] = None
        self._chat_template: Optional[str] = None
        self._chat_template_override: Optional[str] = None
        self._supports_reasoning: bool = False
        self._reasoning_always_on: bool = False
        self._reasoning_style: str = "enable_thinking"
        self._supports_preserve_thinking: bool = False
        self._supports_tools: bool = False
        self._cache_type_kv: Optional[str] = None
        self._reasoning_default: bool = True
        self._speculative_type: Optional[str] = None
        # Canonical UI-facing mode the user requested
        # (auto/mtp/ngram/mtp+ngram/off/ngram-simple). Round-tripped through the
        # status API so the dropdown reflects the picked mode, not the resolved
        # flag set (auto on a 27B MTP GGUF resolves to draft-mtp but reads "Auto").
        self._requested_spec_mode: Optional[str] = None
        # User --spec-draft-n-max override (None = platform default).
        self._spec_draft_n_max: Optional[int] = None
        # KV-cache estimation fields (populated by _read_gguf_metadata)
        self._n_layers: Optional[int] = None
        self._n_kv_heads: Optional[int] = None
        self._n_kv_heads_by_layer: Optional[list[int]] = None
        self._n_heads: Optional[int] = None
        self._embedding_length: Optional[int] = None
        # Architecture-aware KV fields for 5-path estimation
        self._kv_key_length: Optional[int] = None
        self._kv_value_length: Optional[int] = None
        self._sliding_window: Optional[int] = None
        self._sliding_window_pattern: Optional[list[bool]] = None
        self._full_attention_interval: Optional[int] = None
        self._kv_lora_rank: Optional[int] = None
        self._key_length_mla: Optional[int] = None
        self._kv_key_length_swa: Optional[int] = None
        self._kv_value_length_swa: Optional[int] = None
        self._ssm_inner_size: Optional[int] = None
        self._ssm_state_size: Optional[int] = None
        # Last N layers reuse earlier layers' KV and don't allocate their own
        # cache (Gemma 3n / Gemma 4: <arch>.attention.shared_kv_layers).
        self._shared_kv_layers: Optional[int] = None
        # MTP head count (llama.cpp #22673); >0 enables --spec-type draft-mtp.
        self._nextn_predict_layers: Optional[int] = None
        self._lock = threading.Lock()
        # Wraps load_model() end-to-end so concurrent loads serialise and never
        # coexist as two llama-server processes (#5401).
        self._serial_load_lock = threading.Lock()
        # Last extra_args / requested n_ctx, preserved across unload so the chat
        # UI's /unload+/load Apply path can inherit them (#5401).
        # ``_extra_args_source`` records the (model_identifier, hf_variant) the
        # stored args came from so the route can refuse cross-model inheritance.
        self._extra_args: Optional[List[str]] = None
        self._extra_args_source: Optional[tuple[str, Optional[str]]] = None
        self._requested_n_ctx: int = 0
        self._stdout_lines: list[str] = []
        self._stdout_thread: Optional[threading.Thread] = None
        # llama-server tee log (see _drain_stdout / _kill_process).
        self._llama_log_fh = None
        self._llama_log_path: Optional[Path] = None
        self._cancel_event = threading.Event()
        self._api_key: Optional[str] = None
        # True once a probe has completed; cleared on transient failure.
        self._is_audio: bool = False
        self._audio_type: Optional[str] = None
        self._audio_probed: bool = False
        # Audio INPUT capability (distinct from _is_audio, which is TTS output).
        self._has_audio_input: bool = False
        self._mmproj_has_audio: bool = False  # clip.has_audio_encoder, set at load
        # Monotonic timestamp set in _kill_process; read by load_model
        # to decide whether to wait for the VRAM reclaim to finish.
        self._last_kill_monotonic: float = 0.0

        self._kill_orphaned_servers()
        atexit.register(self._cleanup)

    # ── Properties ────────────────────────────────────────────────

    @property
    def is_loaded(self) -> bool:
        return self._process is not None and self._healthy

    @property
    def is_active(self) -> bool:
        """True if a llama-server process exists (loading or loaded)."""
        return self._process is not None

    @property
    def base_url(self) -> str:
        return f"http://127.0.0.1:{self._port}"

    @property
    def model_identifier(self) -> Optional[str]:
        return self._model_identifier

    @property
    def is_vision(self) -> bool:
        return self._is_vision

    @property
    def hf_variant(self) -> Optional[str]:
        return self._hf_variant

    @property
    def extra_args(self) -> Optional[List[str]]:
        """Extra llama-server flags from the last load (a copy). None =
        never set, [] = explicitly cleared. Used by the route for
        inheritance."""
        return list(self._extra_args) if self._extra_args is not None else None

    @property
    def requested_n_ctx(self) -> int:
        """n_ctx the last load was invoked with (not the effective cap).
        0 means Auto. Used by the route to detect Auto-vs-explicit flips."""
        return self._requested_n_ctx

    @property
    def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]:
        """(model_identifier, hf_variant) the stored extra_args came from.
        ``None`` if no extras have ever been recorded. Used by the route
        to refuse cross-model inheritance (#5401)."""
        return self._extra_args_source

    @property
    def context_length(self) -> Optional[int]:
        """Return the effective context length the server is running at."""
        return self._effective_context_length or self._context_length

    @property
    def max_context_length(self) -> Optional[int]:
        """Return the largest context that fits on this hardware at load time.

        The UI's "safe zone" warning threshold: the ``_fit_context_to_vram``
        binary-search cap for the best GPU subset, or the 4096 fallback if the
        weights exceed 90% of every subset. The slider ceiling is
        ``native_context_length``; dragging above this triggers the warning.
        """
        return self._max_context_length or self._context_length

    @property
    def native_context_length(self) -> Optional[int]:
        """Return the model's native context length from GGUF metadata."""
        return self._context_length

    def load_progress(self) -> Optional[dict]:
        """Return live model-load progress, or None if not loading.

        During warm-up llama-server mmaps weight shards into page cache before
        pushing layers to VRAM, a window where status only reports ``loading``
        and the UI spinner looks stuck for minutes on large MoEs. Samples
        ``/proc/<pid>/status VmRSS`` against the sum of GGUF shard sizes for a
        real progress bar. Returns ``None`` when no load is in flight.

        Shape::

            {
                "phase": "mmap" | "ready",
                "bytes_loaded": int,   # VmRSS of the llama-server
                "bytes_total":  int,   # sum of shard file sizes
                "fraction": float,     # bytes_loaded / bytes_total, 0..1
            }

        Linux-only; returns ``None`` where ``/proc/<pid>/status`` is unavailable.
        """
        proc = self._process
        if proc is None:
            return None
        pid = proc.pid
        if pid is None:
            return None

        # Sum shard sizes (primary + any extras alongside).
        bytes_total = 0
        gguf_path = self._gguf_path
        if gguf_path:
            primary = Path(gguf_path)
            try:
                if primary.is_file():
                    bytes_total += primary.stat().st_size
            except OSError:
                pass
            # Extra shards share the primary's prefix before the shard index.
            try:
                parent = primary.parent
                stem = primary.name
                m = _SHARD_RE.match(stem)
                prefix = m.group(1) if m else None
                if prefix and parent.is_dir():
                    for sibling in parent.iterdir():
                        if (
                            sibling.is_file()
                            and sibling.name.startswith(prefix)
                            and sibling.name != stem
                            and sibling.suffix == ".gguf"
                        ):
                            try:
                                bytes_total += sibling.stat().st_size
                            except OSError:
                                pass
            except OSError:
                pass

        # Read VmRSS from /proc/<pid>/status (kilobytes on Linux).
        bytes_loaded = 0
        try:
            with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
                for line in f:
                    if line.startswith("VmRSS:"):
                        kb = int(line.split()[1])
                        bytes_loaded = kb * 1024
                        break
        except (FileNotFoundError, PermissionError, ValueError, OSError):
            return None

        phase = "ready" if self._healthy else "mmap"
        fraction = 0.0
        if bytes_total > 0:
            fraction = min(1.0, bytes_loaded / bytes_total)
        return {
            "phase": phase,
            "bytes_loaded": bytes_loaded,
            "bytes_total": bytes_total,
            "fraction": round(fraction, 4),
        }

    @property
    def chat_template(self) -> Optional[str]:
        return self._chat_template

    @property
    def chat_template_override(self) -> Optional[str]:
        return self._chat_template_override

    @property
    def supports_reasoning(self) -> bool:
        return self._supports_reasoning

    @property
    def reasoning_always_on(self) -> bool:
        return self._reasoning_always_on

    @property
    def reasoning_style(self) -> str:
        return self._reasoning_style

    @property
    def supports_preserve_thinking(self) -> bool:
        return self._supports_preserve_thinking

    @property
    def reasoning_default(self) -> bool:
        return self._reasoning_default

    def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
        if self._reasoning_style == "reasoning_effort":
            return {"reasoning_effort": "high" if enable_thinking else "low"}
        return {"enable_thinking": enable_thinking}

    def _request_reasoning_kwargs(
        self,
        enable_thinking: Optional[bool],
        reasoning_effort: Optional[str] = None,
        preserve_thinking: Optional[bool] = None,
    ) -> Optional[dict]:
        """Build chat_template_kwargs from per-request reasoning fields.

        Merges the active model's reasoning style (``enable_thinking`` or
        ``reasoning_effort``) plus the independent ``preserve_thinking``
        kwarg when the template supports it.
        """
        kwargs: dict = {}
        # Always-on reasoning models hardcode <think> tags and don't consume
        # enable_thinking / reasoning_effort -- skip.
        if self._supports_reasoning and not self._reasoning_always_on:
            if self._reasoning_style == "reasoning_effort":
                if reasoning_effort in ("low", "medium", "high"):
                    kwargs["reasoning_effort"] = reasoning_effort
                elif enable_thinking is not None:
                    kwargs["reasoning_effort"] = "high" if enable_thinking else "low"
            else:
                if enable_thinking is not None:
                    kwargs["enable_thinking"] = enable_thinking
        if self._supports_preserve_thinking and preserve_thinking is not None:
            kwargs["preserve_thinking"] = preserve_thinking
        return kwargs or None

    @property
    def supports_tools(self) -> bool:
        return self._supports_tools

    @property
    def cache_type_kv(self) -> Optional[str]:
        return self._cache_type_kv

    @property
    def speculative_type(self) -> Optional[str]:
        return self._speculative_type

    @property
    def requested_spec_mode(self) -> Optional[str]:
        """Canonical UI-facing mode the user requested (see field doc)."""
        return self._requested_spec_mode

    @property
    def spec_draft_n_max(self) -> Optional[int]:
        """User --spec-draft-n-max override active on the load, or None when
        the platform default (6 GPU / 3 CPU) is in effect."""
        return self._spec_draft_n_max

    # ── Binary discovery ──────────────────────────────────────────

    @staticmethod
    def _find_llama_server_binary() -> Optional[str]:
        """
        Locate the llama-server binary.

        Search order:
        1.  LLAMA_SERVER_PATH environment variable (direct path to binary)
        1b. UNSLOTH_LLAMA_CPP_PATH env var (custom llama.cpp install dir)
        2.  ~/.unsloth/llama.cpp/llama-server        (make build, root dir)
        3.  ~/.unsloth/llama.cpp/build/bin/llama-server  (cmake build, Linux)
        4.  ~/.unsloth/llama.cpp/build/bin/Release/llama-server.exe  (cmake build, Windows)
        5.  ./llama.cpp/llama-server                 (legacy: make build, root dir)
        6.  ./llama.cpp/build/bin/llama-server        (legacy: cmake in-tree build)
        7.  llama-server on PATH                     (system install)
        8.  ./bin/llama-server                       (legacy: extracted binary)
        """
        binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"

        # 1. Env var: direct path to binary
        env_path = os.environ.get("LLAMA_SERVER_PATH")
        if env_path and Path(env_path).is_file():
            return env_path

        # 1b. UNSLOTH_LLAMA_CPP_PATH: custom llama.cpp install dir
        custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
        if custom_llama_cpp:
            custom_dir = Path(custom_llama_cpp)
            # Root dir (make builds)
            root_bin = custom_dir / binary_name
            if root_bin.is_file():
                return str(root_bin)
            # build/bin/ (cmake on Linux)
            cmake_bin = custom_dir / "build" / "bin" / binary_name
            if cmake_bin.is_file():
                return str(cmake_bin)
            # build/bin/Release/ (cmake on Windows)
            if sys.platform == "win32":
                win_bin = custom_dir / "build" / "bin" / "Release" / binary_name
                if win_bin.is_file():
                    return str(win_bin)

        # 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp;
        # default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio).
        legacy_llama = Path.home() / ".unsloth" / "llama.cpp"
        try:
            from utils.paths.storage_roots import studio_root as _sr  # noqa: WPS433

            _resolved_sr = _sr()
            _legacy_studio = Path.home() / ".unsloth" / "studio"
            try:
                _is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve()
            except (OSError, ValueError):
                _is_legacy = _resolved_sr == _legacy_studio
            if _is_legacy:
                search_roots = [legacy_llama]
            else:
                # _kill_orphaned_servers excludes the legacy root in custom
                # mode; discovery must match so we never spawn a server we
                # then refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled
                # earlier) is the explicit way to share a build across roots.
                search_roots = [_resolved_sr / "llama.cpp"]
        except (ImportError, OSError, ValueError):
            search_roots = [legacy_llama]
        _seen_roots: set[str] = set()
        _unique_roots: list[Path] = []
        for r in search_roots:
            k = str(r)
            if k not in _seen_roots:
                _seen_roots.add(k)
                _unique_roots.append(r)
        for unsloth_home in _unique_roots:
            home_root = unsloth_home / binary_name
            if home_root.is_file():
                return str(home_root)
            home_linux = unsloth_home / "build" / "bin" / binary_name
            if home_linux.is_file():
                return str(home_linux)
            if sys.platform == "win32":
                home_win = unsloth_home / "build" / "bin" / "Release" / binary_name
                if home_win.is_file():
                    return str(home_win)

        # 5-6. Legacy: in-tree build (older setup.sh / setup.ps1)
        project_root = Path(__file__).resolve().parents[4]
        # Root dir (make builds)
        root_path = project_root / "llama.cpp" / binary_name
        if root_path.is_file():
            return str(root_path)
        # build/bin/ (cmake builds)
        build_path = project_root / "llama.cpp" / "build" / "bin" / binary_name
        if build_path.is_file():
            return str(build_path)
        if sys.platform == "win32":
            win_path = project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
            if win_path.is_file():
                return str(win_path)

        # 7. System PATH
        system_path = shutil.which("llama-server")
        if system_path:
            return system_path

        # 8. Legacy: extracted to bin/
        bin_path = project_root / "bin" / binary_name
        if bin_path.is_file():
            return str(bin_path)

        return None

    # ── llama-server capability probe ─────────────────────────────

    # Cached on (path, mtime); `unsloth studio update` bumps mtime.
    _capability_cache: dict[tuple[str, int], dict[str, object]] = {}

    @classmethod
    def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]:
        """Parse `llama-server --help` for feature flags. Returns
        {found, mtp_token, supports_mtp, ngram_mod_flavor,
        supports_ngram_mod, spec_draft_n_max_flag}.

        ``ngram_mod_flavor``: ``"new"`` when the post-rename
        ``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args;
        ``"legacy"`` when only the pre-rename
        ``--spec-ngram-size-n / --draft-min / --draft-max`` are real (the
        rename ships stub removal entries for legacy names, told apart by
        the "argument has been removed" description); ``None`` if neither
        set is usable.

        ``spec_draft_n_max_flag``: the flag the binary accepts --
        ``--spec-draft-n-max`` post-rename, ``--draft-max`` on legacy.
        ``None`` means n_max cannot be set.
        """
        bin_path = binary or cls._find_llama_server_binary()
        if not bin_path or not Path(bin_path).is_file():
            return {
                "found": False,
                "mtp_token": None,
                "supports_mtp": False,
                "ngram_mod_flavor": None,
                "supports_ngram_mod": False,
                "spec_draft_n_max_flag": None,
            }
        try:
            mtime = int(Path(bin_path).stat().st_mtime)
        except OSError:
            mtime = 0
        cache_key = (bin_path, mtime)
        cached = cls._capability_cache.get(cache_key)
        if cached is not None:
            return cached

        mtp_token: Optional[str] = None
        ngram_mod_flavor: Optional[str] = None
        spec_draft_n_max_flag: Optional[str] = None
        try:
            result = subprocess.run(
                [bin_path, "--help"],
                capture_output = True,
                text = True,
                timeout = 10,
                check = False,
            )
            help_text = (result.stdout or "") + "\n" + (result.stderr or "")
            # Split into per-flag blocks (each --flag line + its indented
            # continuation), so the "argument has been removed" description
            # sits with its flag.
            blocks: dict[str, str] = {}
            current_flags: list[str] = []
            current_desc: list[str] = []
            for line in help_text.splitlines():
                stripped = line.strip()
                if stripped.startswith("-") and not line.startswith(" "):
                    # New flag line; flush previous.
                    if current_flags:
                        desc = " ".join(current_desc)
                        for f in current_flags:
                            blocks[f] = desc
                    current_flags = []
                    current_desc = [stripped]
                    # Extract long-form flag tokens from the DECLARATION
                    # prefix only (comma-separated aliases). Stop at the
                    # first non-flag token so flag references inside
                    # descriptions are ignored.
                    for tok in re.split(r"[,\s]+", stripped):
                        if tok.startswith("--") and re.match(r"--[A-Za-z][A-Za-z0-9_-]*$", tok):
                            current_flags.append(tok)
                        elif tok.startswith("-") and len(tok) > 1:
                            # short alias like -fa; keep scanning aliases.
                            continue
                        else:
                            # First non-flag token marks end of decl.
                            break
                else:
                    current_desc.append(stripped)
            if current_flags:
                desc = " ".join(current_desc)
                for f in current_flags:
                    blocks[f] = desc

            def _is_real(flag: str) -> bool:
                """True if the flag exists AND is not a removal stub."""
                desc = blocks.get(flag)
                if desc is None:
                    return False
                return "argument has been removed" not in desc

            # MTP token from the --spec-type line.
            spec_line = ""
            for line in help_text.splitlines():
                if "--spec-type" in line:
                    spec_line = line
                    break
            # PR #22673 used draft-mtp; later renamed to mtp.
            if "draft-mtp" in spec_line:
                mtp_token = "draft-mtp"
            elif re.search(r"[|,\[]mtp[|,\]]", spec_line):
                mtp_token = "mtp"

            # ngram-mod flag flavor. Post-rename builds advertise both new
            # args (real) and legacy ones (stubs); pre-rename builds only
            # have legacy ones as real.
            new_ngram_real = (
                _is_real("--spec-ngram-mod-n-match")
                and _is_real("--spec-ngram-mod-n-min")
                and _is_real("--spec-ngram-mod-n-max")
            )
            legacy_ngram_real = (
                _is_real("--spec-ngram-size-n")
                and _is_real("--draft-max")
                and _is_real("--draft-min")
            )
            if new_ngram_real:
                ngram_mod_flavor = "new"
            elif legacy_ngram_real:
                ngram_mod_flavor = "legacy"

            # n_max flag: prefer post-rename, fall back to legacy.
            if _is_real("--spec-draft-n-max"):
                spec_draft_n_max_flag = "--spec-draft-n-max"
            elif _is_real("--draft-max"):
                spec_draft_n_max_flag = "--draft-max"
        except (OSError, subprocess.SubprocessError) as exc:
            logger.debug(f"llama-server --help probe failed: {exc}")

        info = {
            "found": True,
            "mtp_token": mtp_token,
            "supports_mtp": mtp_token is not None,
            "ngram_mod_flavor": ngram_mod_flavor,
            "supports_ngram_mod": ngram_mod_flavor is not None,
            "spec_draft_n_max_flag": spec_draft_n_max_flag,
        }
        cls._capability_cache[cache_key] = info
        return info

    # ── GPU allocation ────────────────────────────────────────────

    @staticmethod
    def _get_gguf_size_bytes(model_path: str) -> int:
        """Total GGUF size in bytes, including split shards."""
        main = Path(model_path)
        total = main.stat().st_size

        # Check for split shards (e.g. model-00001-of-00003.gguf)
        m = _SHARD_FULL_RE.match(main.name)
        if m:
            prefix, _, num_total = m.group(1), m.group(2), m.group(3)
            sibling_pat = re.compile(
                r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$"
            )
            for sibling in main.parent.iterdir():
                if sibling != main and sibling_pat.match(sibling.name):
                    total += sibling.stat().st_size

        return total

    @staticmethod
    def _amd_apu_wants_unified_memory() -> bool:
        """True only for AMD unified-memory APUs (gfx1150/gfx1151), where
        GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
        False elsewhere (the env hurts discrete GPUs). ROCm reuses torch.cuda.*;
        gcnArchName suffix is stripped."""
        try:
            import torch

            if getattr(torch.version, "hip", None) is None:
                return False
            if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
                return False
            for _i in range(torch.cuda.device_count()):
                try:
                    _arch = getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") or ""
                except Exception:
                    continue
                if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
                    return True
        except Exception:
            return False
        return False

    @staticmethod
    def _get_gpu_free_memory() -> list[tuple[int, int]]:
        """Query free memory per GPU.

        Order:
          1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects
             ``CUDA_VISIBLE_DEVICES``.
          2. ``torch.cuda.mem_get_info`` -- universal fallback that works
             on AMD ROCm too (HIP runtime reuses the ``torch.cuda.*``
             namespace). Covers the AMD case for issue #5106 (nvidia-smi
             probe returned [] on AMD) and NVIDIA hosts missing
             ``nvidia-smi`` from PATH.

        Returns list of (gpu_index, free_mib) sorted by index; empty if no
        supported GPU is reachable.
        """
        # ── NVIDIA via nvidia-smi ────────────────────────────────────
        try:
            result = subprocess.run(
                [
                    "nvidia-smi",
                    "--query-gpu=index,memory.free",
                    "--format=csv,noheader,nounits",
                ],
                capture_output = True,
                text = True,
                timeout = 10,
                env = child_env_without_native_path_secret(),
                **_windows_hidden_subprocess_kwargs(),
            )
            if result.returncode == 0:
                allowed: Optional[set[int]] = None
                cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
                if cvd is not None:
                    try:
                        # `if x.strip()` filters trailing-comma masks ("0,1,").
                        # Empty mask (CVD="") yields an empty set -> all GPUs
                        # filtered out, per codebase convention.
                        allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
                    except ValueError:
                        pass
                gpus: list[tuple[int, int]] = []
                for line in result.stdout.strip().splitlines():
                    parts = line.split(",")
                    if len(parts) == 2:
                        idx = int(parts[0].strip())
                        free_mib = int(parts[1].strip())
                        if allowed is not None and idx not in allowed:
                            continue
                        gpus.append((idx, free_mib))
                # Match the docstring's sort-by-id guarantee (driver order isn't).
                gpus.sort(key = lambda g: g[0])
                if gpus:
                    return gpus
        except Exception as e:
            logger.debug(f"nvidia-smi probe failed: {e}")

        # ── Torch fallback (covers AMD ROCm and missing nvidia-smi) ──
        try:
            import torch

            if not hasattr(torch, "cuda") or not torch.cuda.is_available():
                return []
            if not hasattr(torch.cuda, "mem_get_info"):
                return []
            # torch.cuda enumerates GPUs RELATIVE to the visibility mask. We
            # feed these IDs back into the subprocess as CVD, so visible ordinals
            # must be translated to physical indices first; otherwise CVD=2,3
            # gets rewritten to 0,1 and targets the wrong GPUs.
            physical_ids: Optional[list[int]] = None
            # Match utils/hardware/hardware.py::_get_parent_visible_gpu_spec:
            # treat an empty mask (HIP_VISIBLE_DEVICES="") as "no GPUs" rather
            # than falling through. ``or`` would coerce "" to the wrong source.
            if getattr(torch.version, "hip", None) is not None:
                hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
                rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
                cvd = (
                    hip_v
                    if hip_v is not None
                    else rocr_v
                    if rocr_v is not None
                    else os.environ.get("CUDA_VISIBLE_DEVICES")
                )
            else:
                cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
            if cvd is not None:
                try:
                    # Empty mask (CVD="") yields an empty list -> no GPUs,
                    # consistent with the nvidia-smi path.
                    physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
                except ValueError:
                    physical_ids = None
            gpus = []
            for ordinal in range(torch.cuda.device_count()):
                free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal)
                idx = (
                    physical_ids[ordinal]
                    if physical_ids is not None and ordinal < len(physical_ids)
                    else ordinal
                )
                gpus.append((idx, free_bytes // (1024 * 1024)))
            # Match the nvidia-smi path's docstring guarantee of sorted-by-id.
            return sorted(gpus, key = lambda g: g[0])
        except Exception as e:
            logger.debug(f"torch GPU probe failed: {e}")
            return []

    # Skip the wait when the last kill is older than this; the driver has
    # already reclaimed the prior process's allocations.
    _VRAM_SETTLE_WINDOW_S: float = 15.0

    @staticmethod
    def _wait_for_vram_settle(
        max_wait: float = 2.0,
        interval: float = 0.25,
        tolerance_mib: int = 256,
        since_kill: float = 0.0,
    ) -> None:
        """Poll ``_get_gpu_free_memory`` until free VRAM stabilises.

        The driver reclaims a dead process's allocations asynchronously, so
        sampling free memory in the kill-to-spawn window reads artificially low
        and pushes GPU selection toward needless CPU offload (the Apply-reload
        OOM bare-shell launches never see).

        Short-circuits on cold start, stale kill (older than
        ``_VRAM_SETTLE_WINDOW_S``), CPU-only hosts, probe exceptions, and GPU-set
        changes. ``max_wait`` bounds wall-clock time so a wedged ``nvidia-smi``
        can't extend the reload.
        """
        now = time.monotonic()
        if since_kill <= 0.0:
            return
        if now - since_kill > LlamaCppBackend._VRAM_SETTLE_WINDOW_S:
            return
        deadline = now + max_wait

        def _probe_or_none():
            if time.monotonic() >= deadline:
                return None
            try:
                return LlamaCppBackend._get_gpu_free_memory()
            except Exception:
                return None

        prev = _probe_or_none()
        if prev is None or not prev:
            return
        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return
            # Clip the nap so a near-zero ``max_wait`` is respected.
            time.sleep(min(interval, remaining))
            curr = _probe_or_none()
            if curr is None or not curr or len(curr) != len(prev):
                return
            prev_map = dict(prev)
            stable = True
            for idx, free in curr:
                if idx not in prev_map:
                    stable = False
                    break
                prev_free = prev_map[idx]
                # Adaptive: 2% of the larger sample dominates the 256 MiB floor.
                per_gpu_tol = max(tolerance_mib, int(max(free, prev_free) * 0.02))
                if abs(free - prev_free) >= per_gpu_tol:
                    stable = False
                    break
            if stable:
                return
            prev = curr

    # Free-VRAM fraction at which Studio pins the GPU directly instead of
    # deferring to ``--fit on``. 5% headroom covers CUDA context + compute
    # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106).
    _GPU_PIN_VRAM_FRACTION = 0.95

    @staticmethod
    def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
        """Return DLL dirs from pip-installed CUDA wheels under
        ``<prefix>/Lib/site-packages/`` so llama-server.exe can load
        ``cudart64_X.dll`` / ``cublas64_X.dll`` without a system CUDA toolkit.
        Mirrors the Linux ``nvidia/cu*/lib`` LD_LIBRARY_PATH block, covering the
        Windows wheel layouts seen in the wild:
          * ``nvidia/<pkg>/bin`` -- legacy modular wheels.
          * ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- CUDA 13 layout
            for unsuffixed packages (#5106).
          * ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda repacks.
          * ``torch/lib`` -- PyTorch's CUDA-bundled wheel can ship
            ``cudart64_*.dll`` here; mirrors install_llama_prebuilt.py.

        Walks with ``Path.iterdir`` not ``glob.glob`` so it's safe against
        Windows paths containing ``[`` or ``]`` (valid in usernames)."""
        site_packages = Path(prefix) / "Lib" / "site-packages"
        out: list[str] = []
        seen: set[str] = set()

        def _add(path: Path) -> None:
            if not path.is_dir():
                return
            key = os.path.normcase(os.path.abspath(str(path)))
            if key in seen:
                return
            seen.add(key)
            out.append(str(path))

        nvidia_root = site_packages / "nvidia"
        if nvidia_root.is_dir():
            for pkg_dir in nvidia_root.iterdir():
                if not pkg_dir.is_dir():
                    continue
                # Arch-specific subdirs first so the explicit cudart64_X.dll
                # location wins over an empty sibling ``bin``.
                for sub in (
                    pkg_dir / "bin" / "x86_64",
                    pkg_dir / "bin" / "x64",
                    pkg_dir / "bin",
                    pkg_dir / "Library" / "bin" / "x86_64",
                    pkg_dir / "Library" / "bin" / "x64",
                    pkg_dir / "Library" / "bin",
                ):
                    _add(sub)
        _add(site_packages / "torch" / "lib")
        return out

    @staticmethod
    def _build_windows_path_dirs(binary_dir: str, prefix: str, cuda_path: str) -> list[str]:
        """Ordered PATH entries prepended so llama-server.exe resolves cudart /
        cublas DLLs: binary_dir, pip nvidia wheels, CUDA_PATH/bin, .../bin/x64.
        Extracted so test_windows_gpu_detection_mock tests the real logic. #5106."""
        path_dirs = [binary_dir]
        path_dirs.extend(LlamaCppBackend._windows_pip_nvidia_dll_dirs(prefix))
        if cuda_path:
            cuda_bin = os.path.join(cuda_path, "bin")
            if os.path.isdir(cuda_bin):
                path_dirs.append(cuda_bin)
            cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
            if os.path.isdir(cuda_bin_x64):
                path_dirs.append(cuda_bin_x64)
        return path_dirs

    @staticmethod
    def _select_gpus(
        model_size_bytes: int, gpus: list[tuple[int, int]]
    ) -> tuple[Optional[list[int]], bool]:
        """Pick GPU(s) for a model from estimated VRAM and free memory.

        ``model_size_bytes`` should include weights and estimated KV cache.
        ``_GPU_PIN_VRAM_FRACTION`` provides headroom for compute buffers,
        CUDA context, and other runtime overhead.

        Returns (gpu_indices, use_fit):
          - ([1], False)       fits on 1 GPU at the headroom threshold
          - ([1, 2], False)    needs 2 GPUs
          - (None, True)       too large, let --fit handle it
        """
        if not gpus:
            return None, True

        model_size_mib = model_size_bytes / (1024 * 1024)
        usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION

        # Sort GPUs by free memory descending
        ranked = sorted(gpus, key = lambda g: g[1], reverse = True)

        # Try 1 GPU at the usable-VRAM threshold.
        if ranked[0][1] * usable_fraction >= model_size_mib:
            return [ranked[0][0]], False

        # Try N GPUs (accumulate free memory from most-free)
        cumulative = 0
        selected = []
        for idx, free_mib in ranked:
            selected.append(idx)
            cumulative += free_mib * usable_fraction
            if cumulative >= model_size_mib:
                return sorted(selected), False

        # Too large even for all GPUs; let --fit handle it
        logger.debug(
            "Model does not fit in available GPU memory, falling back to --fit",
            model_size_mib = round(model_size_mib, 2),
            ranked_gpus = ranked,
        )
        return None, True

    # ── KV cache VRAM estimation ─────────────────────────────────────

    def _can_estimate_kv(self) -> bool:
        """True if we have enough GGUF metadata to estimate KV cache size."""
        if self._n_layers is None:
            return False
        # MLA: kv_lora_rank suffices (K-only cache).
        if self._kv_lora_rank is not None:
            return True
        # New-style: need explicit key AND value dimensions.
        if self._kv_key_length is not None and self._kv_value_length is not None:
            return True
        # Legacy: need embedding_length + a head count (scalar or per-layer).
        return self._embedding_length is not None and (
            self._n_kv_heads is not None
            or self._n_heads is not None
            or self._n_kv_heads_by_layer is not None
        )

    def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int:
        if self._n_kv_heads_by_layer is not None and layer_idx < len(self._n_kv_heads_by_layer):
            return self._n_kv_heads_by_layer[layer_idx]
        return fallback

    def _estimate_kv_cache_bytes(
        self,
        n_ctx: int,
        cache_type_kv: Optional[str] = None,
        *,
        swa_full: bool = False,
        n_parallel: int = 1,
        kv_unified: bool = True,
        ctx_checkpoints: int = 0,
    ) -> int:
        """Estimate KV cache VRAM for a given context length.

        5-path architecture-aware estimation:
          1. MLA      -- compressed KV latent + RoPE, K-only (no separate V)
          2. Hybrid   -- only attention layers need KV (Mamba layers don't)
          3. SWA      -- sliding-window layers cache min(ctx, window) tokens
          4. GQA      -- standard full KV with explicit key/value dimensions
          5. Legacy   -- fallback using embed // n_heads

        Server-flag knobs (mirror llama-server's CLI):
          swa_full        -- ``--swa-full``: force SWA layers to cache full
                             ``n_ctx`` (collapses path 3 to path 4 for them).
          n_parallel      -- ``--parallel`` slots: non-SWA layers stay constant
                             (cells split across slots), SWA layers scale linearly.
          kv_unified      -- ``--kv-unified`` (default on): no-op for memory math;
                             kept for API forward-compat.
          ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots
                             per slot, one sliding-window of state per SWA layer.

        Returns 0 if metadata is insufficient.
        """
        if not self._can_estimate_kv() or n_ctx <= 0:
            return 0

        n_layers = self._n_layers  # type: ignore[assignment]
        # Gemma 3n / Gemma 4 reuse earlier KV in the last ``shared_kv_layers``
        # blocks (no cache). Floor at 1 so a bad GGUF can't zero out KV.
        shared = self._shared_kv_layers or 0
        n_layers_kv = max(1, n_layers - shared)
        n_kv = self._n_kv_heads or self._n_heads or 1  # type: ignore[assignment]

        # Bytes per element depends on KV cache quantization
        bpe = {
            "f32": 4.0,
            "f16": 2.0,
            "bf16": 2.0,
            "q8_0": 34 / 32,
            "q5_1": 0.75,
            "q5_0": 0.6875,
            "q4_1": 0.625,
            "q4_0": 0.5625,
            "iq4_nl": 0.5625,
        }.get(cache_type_kv or "f16", 2.0)

        slots = max(1, n_parallel)

        # Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
        # One compressed KV latent per token/layer (shared across heads); V is
        # reconstructed from it, no separate V cache. key_length = kv_lora_rank
        # + rope_dim. MLA GGUFs set head_count_kv=1; default to 1 if absent to
        # avoid falling back to n_heads (e.g. 128 for DeepSeek-V3) which 128x's.
        if self._kv_lora_rank is not None:
            n_kv_mla = self._n_kv_heads or 1
            rope_dim = self._key_length_mla or 64
            key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim)
            return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe)

        key_len = self._kv_key_length
        val_len = self._kv_value_length

        # Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B)
        # Only 1 in N layers is attention; the rest are Mamba (no KV cache).
        if self._ssm_inner_size is not None and self._full_attention_interval is not None:
            fai = self._full_attention_interval
            n_attn = -(-n_layers // fai) if fai > 0 else n_layers  # ceiling division
            if key_len is not None and val_len is not None:
                return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe)
            head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
            return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe)

        # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...).
        # Pattern filled by the resolver at parse time; if absent, falls through
        # to the legacy 1/4-global heuristic below. Per-layer-type --parallel N
        # accounting (verified against llama-server):
        #   * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT.
        #   * SWA layers:     per-slot cells = 2*sliding_window (capped at n_ctx
        #                     and per_slot_ctx) -> grows LINEARLY in slots.
        # --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds
        # N snapshots per SWA layer per slot.
        if (
            self._sliding_window is not None
            and self._sliding_window > 0
            and key_len is not None
            and val_len is not None
        ):
            swa = self._sliding_window
            per_slot_ctx = max(1, n_ctx // slots)
            # --swa-full caches full context like non-SWA (per-slot cells =
            # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA
            # caches 2*sliding_window per slot, clamped at per-slot ctx.
            swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx)
            key_len_swa = self._kv_key_length_swa or key_len
            val_len_swa = self._kv_value_length_swa or val_len
            if self._sliding_window_pattern is not None:
                global_bytes = 0.0  # constant across slots
                swa_bytes_per_slot = 0.0  # multiplied by slots
                checkpoint_extra_per_slot = 0.0
                # Only layers that allocate their own KV; trailing shared layers
                # reuse earlier caches.
                for layer_idx in range(n_layers_kv):
                    layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv)
                    is_swa = (
                        layer_idx < len(self._sliding_window_pattern)
                        and self._sliding_window_pattern[layer_idx]
                    )
                    if is_swa:
                        swa_bytes_per_slot += (
                            swa_cells_per_slot * layer_n_kv * (key_len_swa + val_len_swa) * bpe
                        )
                        if ctx_checkpoints > 0 and not swa_full:
                            checkpoint_extra_per_slot += (
                                ctx_checkpoints
                                * swa
                                * layer_n_kv
                                * (key_len_swa + val_len_swa)
                                * bpe
                            )
                    else:
                        global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe
                return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot))
            n_global = max(1, n_layers_kv // 4)
            n_swa = n_layers_kv - n_global
            kv_per_token = n_kv * (key_len + val_len) * bpe
            kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe
            global_bytes = n_global * n_ctx * kv_per_token
            swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa
            checkpoint_extra_per_slot = (
                ctx_checkpoints * n_swa * swa * kv_per_token_swa
                if ctx_checkpoints > 0 and not swa_full
                else 0.0
            )
            return int(global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot))

        # Path 4: Standard GQA with explicit key/value dimensions
        if key_len is not None and val_len is not None:
            return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe)

        # Path 5: Legacy fallback (old GGUFs without explicit dimensions)
        head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
        return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)

    def _fit_context_to_vram(
        self,
        requested_ctx: int,
        available_mib: int,
        model_size_bytes: int,
        cache_type_kv: Optional[str] = None,
        min_ctx: int = 4096,
        *,
        swa_full: bool = False,
        n_parallel: int = 1,
        kv_unified: bool = True,
        ctx_checkpoints: int = 0,
        kv_on_gpu: bool = True,
        mtp_engaged: bool = False,
    ) -> int:
        """Return the largest context length that fits in GPU VRAM.

        Uses 90% of available VRAM as the ctx-fit budget -- tighter than
        ``_GPU_PIN_VRAM_FRACTION`` on purpose (over-promising context OOMs at
        runtime). If the weights alone don't fit, returns ``requested_ctx``.

        ``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV
        cache lives in CPU RAM and the requested context is honored verbatim.
        Other keyword args mirror ``_estimate_kv_cache_bytes``.

        ``mtp_engaged`` reserves extra VRAM for the MTP draft model's KV cache +
        compute buffers, else tight tiers (e.g. 32 GB) spill to a slower path.
        """
        if not self._can_estimate_kv():
            logger.debug(
                "Skipping context fit because KV cache metadata is unavailable",
                requested_ctx = requested_ctx,
                available_mib = available_mib,
            )
            return requested_ctx

        # KV lives off-GPU: no VRAM accounting needed for the cache itself.
        if not kv_on_gpu:
            return requested_ctx

        kv_kwargs = dict(
            swa_full = swa_full,
            n_parallel = n_parallel,
            kv_unified = kv_unified,
            ctx_checkpoints = ctx_checkpoints,
        )

        # MTP needs a tighter budget; drop from 0.90 to 0.85.
        budget_frac = 0.85 if mtp_engaged else 0.90
        budget_bytes = available_mib * 1024 * 1024 * budget_frac
        model_footprint = model_size_bytes

        # Already fits?
        kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
        if model_footprint + kv <= budget_bytes:
            return requested_ctx

        # Weights alone exceed budget -- reducing ctx can't help; --fit handles it.
        if model_footprint >= budget_bytes:
            logger.debug(
                "Model footprint exceeds GPU budget before KV cache",
                requested_ctx = requested_ctx,
                available_mib = available_mib,
                model_size_gb = round(model_footprint / (1024**3), 2),
            )
            return requested_ctx

        # Binary search for max context that fits
        remaining = budget_bytes - model_footprint
        effective_min = min(min_ctx, requested_ctx)
        lo, hi = effective_min, requested_ctx
        best = effective_min
        while lo <= hi:
            mid = (lo + hi) // 2
            kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
            if kv <= remaining:
                best = mid
                lo = mid + 1
            else:
                hi = mid - 1

        # Round down to nearest 256 for alignment, never above requested_ctx
        best = (best // 256) * 256
        best = max(effective_min, best)
        best = min(best, requested_ctx)
        return best

    # ── Variant fallback ────────────────────────────────────────────

    @staticmethod
    def _find_smallest_fitting_variant(
        hf_repo: str,
        free_bytes: int,
        hf_token: Optional[str] = None,
    ) -> Optional[tuple[str, int]]:
        """Find the smallest GGUF variant (including all shards) that fits.

        Groups split shards by variant prefix and sums their sizes (e.g.
        UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total).

        Returns (first_shard_filename, total_size_bytes) or None.
        """
        try:
            from huggingface_hub import get_paths_info, list_repo_files

            files = list_repo_files(hf_repo, token = hf_token)
            gguf_files = [f for f in files if f.endswith(".gguf") and "mmproj" not in f.lower()]
            if not gguf_files:
                return None

            # Sizes for all GGUF files
            path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
            size_map = {p.path: (p.size or 0) for p in path_infos}

            # Group by variant: shards share a prefix before -NNNNN-of-NNNNN
            variants: dict[str, list[str]] = {}
            for f in gguf_files:
                m = _SHARD_RE.match(f)
                key = m.group(1) if m else f
                variants.setdefault(key, []).append(f)

            # Sum shard sizes per variant, track the first shard (for download)
            variant_sizes: list[tuple[str, int, list[str]]] = []
            for key, shard_files in variants.items():
                total = sum(size_map.get(f, 0) for f in shard_files)
                first = sorted(shard_files)[0]
                variant_sizes.append((first, total, shard_files))

            # Smallest that fits
            variant_sizes.sort(key = lambda x: x[1])
            for first_file, total_size, _ in variant_sizes:
                if total_size > 0 and total_size <= free_bytes:
                    return first_file, total_size

            return None
        except Exception:
            return None

    # ── Port allocation ───────────────────────────────────────────

    @staticmethod
    def _find_free_port() -> int:
        """Find an available TCP port."""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("127.0.0.1", 0))
            return s.getsockname()[1]

    # ── Stdout drain (prevents pipe deadlock on Windows) ─────────

    def _drain_stdout(self):
        """Read subprocess stdout lines in a background thread.

        Prevents a pipe-buffer deadlock on Windows (~4 KB buffer): without
        draining, llama-server blocks on writes and never becomes healthy.
        Each line is also teed to ``self._llama_log_fh`` when set, so a
        post-mortem has the full output even if the crash predates the
        drain-thread join in ``_wait_for_health``.
        """
        try:
            for line in self._process.stdout:
                line = line.rstrip()
                if line:
                    self._stdout_lines.append(line)
                    logger.debug(f"[llama-server] {line}")
                    fh = getattr(self, "_llama_log_fh", None)
                    if fh is not None:
                        try:
                            fh.write(line + "\n")
                            fh.flush()
                        except (ValueError, OSError):
                            # Log file closed under us; tee silently.
                            pass
        except (ValueError, OSError):
            # Pipe closed -- process terminating.
            pass

    # GGUF KV type sizes for fast skipping
    _GGUF_TYPE_SIZE = {
        0: 1,
        1: 1,
        2: 2,
        3: 2,
        4: 4,
        5: 4,
        6: 4,
        7: 1,
        10: 8,
        11: 8,
        12: 8,
    }

    @staticmethod
    def _gguf_skip_value(f, vtype: int) -> None:
        """Skip a GGUF KV value without reading it."""
        sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(vtype)
        if sz is not None:
            f.seek(sz, 1)
        elif vtype == 8:  # STRING
            slen = struct.unpack("<Q", f.read(8))[0]
            f.seek(slen, 1)
        elif vtype == 9:  # ARRAY
            atype = struct.unpack("<I", f.read(4))[0]
            alen = struct.unpack("<Q", f.read(8))[0]
            elem_sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(atype)
            if elem_sz is not None:
                f.seek(elem_sz * alen, 1)
            elif atype == 8:
                for _ in range(alen):
                    slen = struct.unpack("<Q", f.read(8))[0]
                    f.seek(slen, 1)
            else:
                for _ in range(alen):
                    LlamaCppBackend._gguf_skip_value(f, atype)

    @staticmethod
    def _gguf_read_array_value(f, atype: int, alen: int) -> Optional[list]:
        if atype == 4:  # UINT32
            return [struct.unpack("<I", f.read(4))[0] for _ in range(alen)]
        if atype == 5:  # INT32
            return [struct.unpack("<i", f.read(4))[0] for _ in range(alen)]
        if atype == 7:  # BOOL
            return [struct.unpack("<?", f.read(1))[0] for _ in range(alen)]

        for _ in range(alen):
            LlamaCppBackend._gguf_skip_value(f, atype)
        return None

    def _read_gguf_metadata(self, gguf_path: str) -> None:
        """Read context_length, architecture params, and chat_template from a GGUF header.

        Parses only the KV pairs we need (~30ms even for multi-GB files).
        For split GGUFs, metadata is always in shard 1.
        """
        # Reset metadata so stale flags (e.g. _supports_reasoning) don't
        # carry over when switching models.
        self._context_length = None
        self._chat_template = None
        self._supports_reasoning = False
        self._reasoning_always_on = False
        self._reasoning_style = "enable_thinking"
        self._reasoning_default = True
        self._supports_preserve_thinking = False
        self._supports_tools = False
        self._n_layers = None
        self._n_kv_heads = None
        self._n_kv_heads_by_layer = None
        self._n_heads = None
        self._embedding_length = None
        self._kv_key_length = None
        self._kv_value_length = None
        self._sliding_window = None
        self._sliding_window_pattern = None
        self._full_attention_interval = None
        self._kv_lora_rank = None
        self._key_length_mla = None
        self._kv_key_length_swa = None
        self._kv_value_length_swa = None
        self._ssm_inner_size = None
        self._ssm_state_size = None
        self._shared_kv_layers = None
        self._nextn_predict_layers = None

        try:
            WANTED = {
                "general.architecture",
                "tokenizer.chat_template",
                # Source-repo hints for the SWA resolver's HF fallback.
                "general.source.huggingface.repository",
                "general.source.url",
                "general.source.repo_url",
                "general.base_model.0.repo_url",
                "general.base_model.0.organization",
                "general.base_model.0.name",
                "general.basename",
                "general.organization",
                "general.size_label",
                "general.finetune",
            }
            # Arch-specific keys added dynamically once we know the arch.
            arch_keys: dict[str, str] = {}  # gguf_key -> attribute name
            arch = None
            sliding_window_pattern_period: Optional[int] = None
            general: dict[str, str] = {}

            with open(gguf_path, "rb") as f:
                magic = struct.unpack("<I", f.read(4))[0]
                if magic != 0x46554747:  # b"GGUF" as little-endian u32
                    return
                _version = struct.unpack("<I", f.read(4))[0]
                _tensor_count, kv_count = struct.unpack("<QQ", f.read(16))

                for _ in range(kv_count):
                    # Tolerate truncated input (e.g. a partial header from an
                    # HTTP byte-range fetch): bail out so the resolver
                    # fallback runs on whatever we parsed.
                    try:
                        key_len_bytes = f.read(8)
                        if len(key_len_bytes) < 8:
                            break
                        key_len = struct.unpack("<Q", key_len_bytes)[0]
                        key_bytes = f.read(key_len)
                        if len(key_bytes) < key_len:
                            break
                        key = key_bytes.decode("utf-8")
                        vtype_bytes = f.read(4)
                        if len(vtype_bytes) < 4:
                            break
                        vtype = struct.unpack("<I", vtype_bytes)[0]
                    except (struct.error, UnicodeDecodeError):
                        break

                    try:
                        if key in WANTED or key in arch_keys:
                            if vtype == 8:  # STRING
                                slen = struct.unpack("<Q", f.read(8))[0]
                                val_s = f.read(slen).decode("utf-8")
                                if key.startswith("general.") and key != "general.architecture":
                                    general[key] = val_s
                                if key == "general.architecture":
                                    arch = val_s
                                    arch_keys = {
                                        f"{arch}.context_length": "context_length",
                                        f"{arch}.block_count": "n_layers",
                                        f"{arch}.attention.head_count_kv": "n_kv_heads",
                                        f"{arch}.attention.head_count": "n_heads",
                                        f"{arch}.embedding_length": "embedding_length",
                                        f"{arch}.attention.key_length": "kv_key_length",
                                        f"{arch}.attention.value_length": "kv_value_length",
                                        f"{arch}.attention.sliding_window": "sliding_window",
                                        f"{arch}.attention.sliding_window_pattern": "sliding_window_pattern",
                                        f"{arch}.full_attention_interval": "full_attention_interval",
                                        f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
                                        f"{arch}.attention.key_length_mla": "key_length_mla",
                                        f"{arch}.attention.key_length_swa": "kv_key_length_swa",
                                        f"{arch}.attention.value_length_swa": "kv_value_length_swa",
                                        f"{arch}.attention.shared_kv_layers": "shared_kv_layers",
                                        f"{arch}.ssm.inner_size": "ssm_inner_size",
                                        f"{arch}.ssm.state_size": "ssm_state_size",
                                        f"{arch}.nextn_predict_layers": "nextn_predict_layers",
                                    }
                                elif key == "tokenizer.chat_template":
                                    self._chat_template = val_s
                            elif vtype in (4, 10):  # UINT32 or UINT64
                                val_i = (
                                    struct.unpack("<I", f.read(4))[0]
                                    if vtype == 4
                                    else struct.unpack("<Q", f.read(8))[0]
                                )
                                attr = arch_keys.get(key)
                                if attr:
                                    if attr == "sliding_window_pattern":
                                        sliding_window_pattern_period = val_i
                                    else:
                                        setattr(self, f"_{attr}", val_i)
                            elif vtype == 9:  # ARRAY
                                atype = struct.unpack("<I", f.read(4))[0]
                                alen = struct.unpack("<Q", f.read(8))[0]
                                val_a = self._gguf_read_array_value(f, atype, alen)
                                attr = arch_keys.get(key)
                                if attr == "n_kv_heads" and val_a is not None:
                                    self._n_kv_heads_by_layer = [int(x) for x in val_a]
                                    if self._n_kv_heads is None and val_a:
                                        self._n_kv_heads = max(int(x) for x in val_a)
                                elif attr == "sliding_window_pattern" and val_a is not None:
                                    self._sliding_window_pattern = [bool(x) for x in val_a]
                                    sliding_window_pattern_period = None
                            else:
                                self._gguf_skip_value(f, vtype)
                        else:
                            self._gguf_skip_value(f, vtype)
                    except (struct.error, UnicodeDecodeError):
                        # Truncated input (e.g. HTTP byte-range header
                        # fetch); break so the resolver fallback runs on
                        # what we have.
                        break

            # Expand a scalar period straight from the GGUF first.
            if (
                self._sliding_window_pattern is None
                and sliding_window_pattern_period
                and self._n_layers
            ):
                self._sliding_window_pattern = [
                    (i + 1) % sliding_window_pattern_period != 0 for i in range(self._n_layers)
                ]

            # Otherwise hand off to the resolver (cache / bootstrap /
            # transformers / HF); see `_resolve_swa_pattern`.
            if self._sliding_window_pattern is None and self._sliding_window and self._n_layers:
                hf_repo_candidates = (
                    general.get("general.source.huggingface.repository"),
                    _hf_repo_from_url(general.get("general.source.url")),
                    _hf_repo_from_url(general.get("general.source.repo_url")),
                    _hf_repo_from_url(general.get("general.base_model.0.repo_url")),
                    (
                        f"{general['general.base_model.0.organization']}/"
                        f"{general['general.base_model.0.name']}".replace(" ", "-")
                        if general.get("general.base_model.0.organization")
                        and general.get("general.base_model.0.name")
                        else None
                    ),
                    (
                        f"{general['general.organization']}/"
                        f"{general['general.basename']}".replace(" ", "-")
                        if general.get("general.organization") and general.get("general.basename")
                        else None
                    ),
                )
                self._sliding_window_pattern = _resolve_swa_pattern(
                    arch,
                    self._n_layers,
                    hf_repo_candidates,
                )

            if self._context_length:
                logger.info(f"GGUF metadata: context_length={self._context_length}")
            if self._chat_template:
                logger.info(f"GGUF metadata: chat_template={len(self._chat_template)} chars")
                # Detect thinking/reasoning support from chat template.
                flags = detect_reasoning_flags(
                    self._chat_template,
                    self._model_identifier,
                    log_source = "GGUF metadata",
                )
                self._supports_reasoning = flags["supports_reasoning"]
                self._reasoning_style = flags["reasoning_style"]
                self._reasoning_always_on = flags["reasoning_always_on"]
                self._supports_preserve_thinking = flags["supports_preserve_thinking"]
                self._supports_tools = flags["supports_tools"]
        except Exception as e:
            logger.warning(f"Failed to read GGUF metadata: {e}")

    # ── HF download (no lock held) ───────────────────────────────

    def _download_gguf(
        self,
        *,
        hf_repo: str,
        hf_variant: Optional[str] = None,
        hf_token: Optional[str] = None,
    ) -> str:
        """Download GGUF file(s) from HuggingFace. Returns local path.

        Runs WITHOUT self._lock so unload_model() can set _cancel_event at
        any time; checks it between each shard download.
        """
        try:
            from huggingface_hub import hf_hub_download
        except ImportError:
            raise RuntimeError(
                "huggingface_hub is required for HF model loading. "
                "Install it with: pip install huggingface_hub"
            )

        # Resolve the filename from the variant
        gguf_filename = None
        gguf_extra_shards: list[str] = []
        if hf_variant:
            try:
                from huggingface_hub import list_repo_files

                files = list_repo_files(hf_repo, token = hf_token)
                variant_lower = hf_variant.lower()
                boundary = re.compile(
                    r"(?<![a-zA-Z0-9])" + re.escape(variant_lower) + r"(?![a-zA-Z0-9])"
                )
                gguf_files = sorted(
                    f for f in files if f.endswith(".gguf") and boundary.search(f.lower())
                )
                if gguf_files:
                    gguf_filename = gguf_files[0]
                    m = _SHARD_FULL_RE.match(gguf_filename)
                    if m:
                        prefix = m.group(1)
                        total = m.group(3)
                        sibling_pat = re.compile(
                            r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$"
                        )
                        gguf_extra_shards = [f for f in gguf_files[1:] if sibling_pat.match(f)]
            except Exception as e:
                logger.warning(f"Could not list repo files: {e}")

            # Offline: resolve variant -> filename from the local HF cache.
            # The heuristic below assumes filenames echo the repo name, which
            # breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match
            # against the rel path (not just basename) so subdir layouts like
            # ``BF16/foo.gguf`` are findable.
            if not gguf_filename:
                try:
                    from utils.models.model_config import _iter_hf_cache_snapshots
                    boundary = re.compile(
                        r"(?<![a-zA-Z0-9])" + re.escape(hf_variant.lower()) + r"(?![a-zA-Z0-9])"
                    )
                    for snap in _iter_hf_cache_snapshots(hf_repo):
                        matches = sorted(
                            p.relative_to(snap).as_posix()
                            for p in snap.rglob("*.gguf")
                            if "mmproj" not in p.name.lower()
                            and boundary.search(p.relative_to(snap).as_posix().lower())
                        )
                        if not matches:
                            continue
                        gguf_filename = matches[0]
                        m = _SHARD_FULL_RE.match(Path(gguf_filename).name)
                        if m:
                            prefix = m.group(1)
                            total = m.group(3)
                            sibling_pat = re.compile(
                                r"^"
                                + re.escape(prefix)
                                + r"-\d{5}-of-"
                                + re.escape(total)
                                + r"\.gguf$"
                            )
                            gguf_extra_shards = [
                                f for f in matches[1:] if sibling_pat.match(Path(f).name)
                            ]
                        logger.info(
                            "Resolved variant %s -> %s from local HF cache",
                            hf_variant,
                            gguf_filename,
                        )
                        break
                except Exception as e:
                    logger.debug(f"Offline cache lookup for variant failed: {e}")

            if not gguf_filename:
                repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
                gguf_filename = f"{repo_name}-{hf_variant}.gguf"

        # Check disk space; fall back to a smaller variant if needed
        all_gguf_files = [gguf_filename] + gguf_extra_shards
        try:
            from huggingface_hub import get_paths_info, try_to_load_from_cache

            path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token))
            total_bytes = sum((p.size or 0) for p in path_infos)

            # Subtract bytes already in the HF cache so we only preflight
            # against what we must download. Without this, re-loading a
            # cached large model (e.g. MiniMax-M2.7-GGUF at 131 GB) fails
            # cold whenever free disk is below the full weight footprint,
            # even though nothing needs downloading.
            already_cached_bytes = 0
            for p in path_infos:
                if not p.size:
                    continue
                try:
                    cached_path = try_to_load_from_cache(hf_repo, p.path)
                except Exception:
                    cached_path = None
                if isinstance(cached_path, str) and os.path.exists(cached_path):
                    try:
                        on_disk = os.path.getsize(cached_path)
                    except OSError:
                        on_disk = 0
                    # Satisfied only when the full blob is present.
                    if on_disk >= p.size:
                        already_cached_bytes += p.size

            total_download_bytes = max(0, total_bytes - already_cached_bytes)

            if total_download_bytes > 0:
                cache_dir = os.environ.get(
                    "HF_HUB_CACHE",
                    str(Path.home() / ".cache" / "huggingface" / "hub"),
                )
                Path(cache_dir).mkdir(parents = True, exist_ok = True)
                free_bytes = shutil.disk_usage(cache_dir).free

                total_gb = total_download_bytes / (1024**3)
                free_gb = free_bytes / (1024**3)
                cached_gb = already_cached_bytes / (1024**3)

                logger.info(
                    f"GGUF download: {total_gb:.1f} GB needed "
                    f"({cached_gb:.1f} GB already cached), "
                    f"{free_gb:.1f} GB free on disk"
                )

                if total_download_bytes > free_bytes:
                    smaller = self._find_smallest_fitting_variant(
                        hf_repo,
                        free_bytes,
                        hf_token,
                    )
                    if smaller:
                        fallback_file, fallback_size = smaller
                        logger.info(
                            f"Selected variant too large ({total_gb:.1f} GB), "
                            f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
                        )
                        gguf_filename = fallback_file
                        _m = _SHARD_RE.match(gguf_filename)
                        _prefix = _m.group(1) if _m else None
                        if _prefix:
                            gguf_extra_shards = sorted(
                                f
                                for f in all_gguf_files
                                if f.startswith(_prefix)
                                and f != gguf_filename
                                and "mmproj" not in f.lower()
                            )
                        else:
                            gguf_extra_shards = []
                    else:
                        raise RuntimeError(
                            f"Not enough disk space to download any variant. "
                            f"Only {free_gb:.1f} GB free in {cache_dir}"
                        )
        except RuntimeError:
            raise
        except Exception as e:
            logger.warning(f"Could not check disk space: {e}")

        gguf_label = f"{hf_repo}/{gguf_filename}" + (
            f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""
        )
        logger.info(f"Resolving GGUF: {gguf_label}")
        try:
            if self._cancel_event.is_set():
                raise RuntimeError("Cancelled")
            dl_start = time.monotonic()
            local_path = hf_hub_download(
                repo_id = hf_repo,
                filename = gguf_filename,
                token = hf_token,
            )
            for shard in gguf_extra_shards:
                if self._cancel_event.is_set():
                    raise RuntimeError("Cancelled")
                logger.info(f"Resolving GGUF shard: {shard}")
                hf_hub_download(
                    repo_id = hf_repo,
                    filename = shard,
                    token = hf_token,
                )
        except RuntimeError as e:
            if "Cancelled" in str(e):
                raise
            raise RuntimeError(
                f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
            )
        except Exception as e:
            raise RuntimeError(
                f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
            )

        dl_elapsed = time.monotonic() - dl_start
        if dl_elapsed < 2.0:
            logger.info(f"GGUF resolved from cache: {local_path}")
        else:
            logger.info(f"GGUF downloaded in {dl_elapsed:.1f}s: {local_path}")
        return local_path

    def _download_mmproj(
        self,
        *,
        hf_repo: str,
        hf_token: Optional[str] = None,
    ) -> Optional[str]:
        """Download the mmproj (vision projection) file from a GGUF repo.

        Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local
        path, or None if none exists.
        """

        def _pick_mmproj(candidates: list[str]) -> Optional[str]:
            mmproj_files = sorted(
                f
                for f in candidates
                if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower()
            )
            if not mmproj_files:
                return None
            for f in mmproj_files:
                if f.lower().endswith("-f16.gguf"):
                    return f
            return mmproj_files[0]

        target: Optional[str] = None
        try:
            from huggingface_hub import list_repo_files
            target = _pick_mmproj(list_repo_files(hf_repo, token = hf_token))
        except Exception as e:
            logger.debug(f"Could not list repo files for mmproj: {e}")

        # Offline: resolve mmproj from the local HF cache snapshot, same as
        # _download_gguf's offline fallback above.
        if target is None:
            try:
                from utils.models.model_config import _iter_hf_cache_snapshots
                for snap in _iter_hf_cache_snapshots(hf_repo):
                    rel_files = [p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")]
                    target = _pick_mmproj(rel_files)
                    if target is not None:
                        logger.info("Resolved mmproj %s from local HF cache", target)
                        break
            except Exception as e:
                logger.debug(f"Offline cache lookup for mmproj failed: {e}")

        if target is None:
            return None

        try:
            from huggingface_hub import hf_hub_download

            logger.info(f"Downloading mmproj: {hf_repo}/{target}")
            local_path = hf_hub_download(
                repo_id = hf_repo,
                filename = target,
                token = hf_token,
            )
            return local_path
        except Exception as e:
            logger.warning(f"Could not download mmproj: {e}")
            return None

    def _resolve_launch_mmproj_path(
        self, *, model_path: str, mmproj_path: Optional[str]
    ) -> Optional[str]:
        """Return mmproj_path iff it exists on disk AND matches the model family.

        None if mmproj_path is None, missing, or family-mismatched.
        """
        if not mmproj_path:
            return None

        mmproj = Path(mmproj_path)
        if not mmproj.is_file():
            logger.warning(f"mmproj file not found: {mmproj_path}")
            return None

        from utils.models.model_config import mmproj_matches_model_family

        if not mmproj_matches_model_family(model_path, str(mmproj)):
            logger.warning(
                f"mmproj does not match model family: model={Path(model_path).name} "
                f"mmproj={mmproj.name}"
            )
            return None

        return str(mmproj)

    # ── Lifecycle ─────────────────────────────────────────────────

    # GGUF ``general.architecture`` values for diffusion / image models.
    # llama.cpp has no such architectures, so loading one as a chat model dies
    # with "unknown model architecture: '<arch>'". These match the patched
    # stable-diffusion.cpp / ComfyUI-GGUF enums. Unsloth publishes FLUX and
    # Qwen-Image GGUFs under
    # https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs.
    # Matched exactly (not a substring) so a chat arch containing "wan"/"sd1"
    # (e.g. "taiwan") isn't misrouted to Images.
    _DIFFUSION_ARCHES = frozenset(
        (
            "qwen_image",
            "flux",
            "sd1",
            "sdxl",
            "sd3",
            "aura",
            "hidream",
            "cosmos",
            "ltxv",
            "hyvid",
            "wan",
            "lumina2",
        )
    )

    @staticmethod
    def _classify_llama_start_failure(
        output: str, gguf_path: Optional[str], model_identifier: Optional[str]
    ) -> str:
        """Explain *why* llama-server failed to start, from its output.

        Several distinct failures otherwise collapse into the same opaque
        "invalid GGUF or out of memory" message. Worst case: a diffusion GGUF
        loaded as a chat model -- valid file, plenty of memory, but llama.cpp
        has no such architecture, so the user is told to free memory that was
        never the problem (#5842). Pick the most specific message we can.
        """
        lowered = (output or "").lower()

        # Detect Ollama source up front so the arch branch can keep the
        # Ollama hint instead of the generic "unsupported arch" message.
        gguf = gguf_path or ""
        is_ollama = (
            ".studio_links" in gguf
            or os.sep + "ollama_links" + os.sep in gguf
            or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf
            or (model_identifier or "").startswith("ollama/")
        )

        # "unknown model architecture: '<arch>'": diffusion -> Images page,
        # Ollama -> Ollama hint, else a precise "unsupported" message. Exact
        # match so chat archs aren't misrouted.
        arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered)
        if arch_match:
            arch = arch_match.group(1)
            if arch in LlamaCppBackend._DIFFUSION_ARCHES:
                return (
                    f"'{arch}' is a diffusion (image-generation) GGUF, which "
                    "llama-server cannot run as a chat/completion model. Use "
                    "Studio's Images page to generate with local diffusion "
                    "GGUFs such as FLUX and Qwen-Image."
                )
            if is_ollama:
                return (
                    "Some Ollama models do not work with llama.cpp. Try a "
                    "different model, or use this model directly through "
                    "Ollama instead."
                )
            return (
                f"llama.cpp does not support this GGUF's model architecture "
                f"('{arch}'). The file is valid, but this model type cannot "
                "be run with llama-server."
            )

        # Other Ollama compat failures that don't name an arch. Only when
        # the output shows a GGUF compat issue, not OOM / missing binaries.
        if is_ollama:
            gguf_compat_hints = (
                "key not found",
                "unknown model architecture",
                "failed to load model",
            )
            if any(h in lowered for h in gguf_compat_hints):
                return (
                    "Some Ollama models do not work with llama.cpp. Try a "
                    "different model, or use this model directly through "
                    "Ollama instead."
                )

        # Fallback: genuinely unknown failure (OOM, missing binary ...).
        return (
            "llama-server failed to start. "
            "Check that the GGUF file is valid and you have enough memory."
        )

    def load_model(
        self,
        *,
        # Local mode: pass a path to a .gguf file
        gguf_path: Optional[str] = None,
        # Vision projection (mmproj) for local vision models
        mmproj_path: Optional[str] = None,
        # HF mode: let llama-server download via -hf "repo:quant"
        hf_repo: Optional[str] = None,
        hf_variant: Optional[str] = None,
        hf_token: Optional[str] = None,
        # Common
        model_identifier: str,
        is_vision: bool = False,
        n_ctx: int = 4096,
        chat_template_override: Optional[str] = None,
        cache_type_kv: Optional[str] = None,
        speculative_type: Optional[str] = None,
        spec_draft_n_max: Optional[int] = None,
        n_threads: Optional[int] = None,
        n_gpu_layers: Optional[int] = None,  # caller compat, unused
        n_parallel: int = 1,
        extra_args: Optional[List[str]] = None,
    ) -> bool:
        """Start llama-server with a GGUF model.

        Two modes:
        - Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
        - HF:    ``hf_repo="...-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``

        Returns True if the server started and the health check passed.
        """
        # Serialise the whole load so concurrent /load calls never leave two
        # llama-server processes alive (#5401 / #5161). Doesn't block /unload.
        with self._serial_load_lock:
            # Duplicate /load that raced past the route check: do nothing if the
            # live server already satisfies this request.
            if self._already_in_target_state(
                gguf_path = gguf_path,
                model_identifier = model_identifier,
                hf_variant = hf_variant,
                n_ctx = n_ctx,
                cache_type_kv = cache_type_kv,
                speculative_type = speculative_type,
                spec_draft_n_max = spec_draft_n_max,
                chat_template_override = chat_template_override,
                extra_args = extra_args,
                is_vision = is_vision,
            ):
                logger.info(
                    f"load_model: backend already in target state for "
                    f"'{model_identifier}', skipping reload"
                )
                # Retry probe only if a prior attempt didn't finish.
                if not self._audio_probed:
                    try:
                        detected = self._detect_audio_type_strict()
                        self._audio_probed = True
                    except Exception as exc:
                        logger.debug("Fast-path audio probe failed: %s", exc)
                        detected = None
                    if detected in ("snac", "bicodec", "dac"):
                        with self._lock:
                            if not self._healthy:
                                return False
                            try:
                                self.init_audio_codec(detected)
                                self._is_audio = True
                                self._audio_type = detected
                            except Exception as exc:
                                logger.warning(
                                    "Failed to init audio codec '%s': %s",
                                    detected,
                                    exc,
                                )
                                self._audio_probed = False
                                return False
                    elif detected:
                        # csm / whisper / audio_vlm: track type but keep
                        # _is_audio False -- GGUF TTS routing only fires for
                        # snac/bicodec/dac.
                        with self._lock:
                            if not self._healthy:
                                return False
                            self._audio_type = detected
                    # Re-derive after a retried probe (_mmproj_has_audio persists).
                    from utils.models.model_config import is_audio_input_type

                    self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool(
                        self._mmproj_has_audio
                    )
                if not self._healthy:
                    return False
                return True

            self._cancel_event.clear()

            # ── Phase 1: kill old process (under lock, fast) ──────────
            with self._lock:
                self._kill_process()

            binary = self._find_llama_server_binary()
            if not binary:
                raise RuntimeError(
                    "llama-server binary not found. "
                    "Run setup.sh to build it, install llama.cpp, "
                    "or set LLAMA_SERVER_PATH environment variable."
                )

            # ── Phase 2: download (NO lock held, so cancel can proceed) ──
            # Scope HF_HUB_OFFLINE to the download block only when DNS is
            # dead; cleanup runs even on exception so a transient hiccup
            # can't quarantine future loads.
            if hf_repo:
                with _hf_offline_if_dns_dead():
                    model_path = self._download_gguf(
                        hf_repo = hf_repo,
                        hf_variant = hf_variant,
                        hf_token = hf_token,
                    )
                    # Auto-download mmproj for vision models
                    if is_vision and not mmproj_path:
                        mmproj_path = self._download_mmproj(
                            hf_repo = hf_repo,
                            hf_token = hf_token,
                        )
            elif gguf_path:
                if not Path(gguf_path).is_file():
                    raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
                model_path = gguf_path
            else:
                raise ValueError("Either gguf_path or hf_repo must be provided")

            # Set identifier early so _read_gguf_metadata can use it (DeepSeek).
            self._model_identifier = model_identifier

            # Read GGUF metadata (context_length, chat_template); header-only.
            self._read_gguf_metadata(model_path)

            if self._cancel_event.is_set():
                logger.info("Load cancelled after download phase")
                return False

            # Outside ``self._lock`` so /unload, /cancel, /status aren't
            # blocked. ``unload_model`` also records the kill, so the
            # frontend /unload+/load Apply path engages the wait here even
            # without an in-process kill.
            self._wait_for_vram_settle(since_kill = self._last_kill_monotonic)

            # ── Phase 3: start llama-server (under lock) ──────────────
            with self._lock:
                # Re-check cancel inside lock
                if self._cancel_event.is_set():
                    logger.info("Load cancelled before server start")
                    return False

                self._port = self._find_free_port()

                # Select GPU(s) from model size + estimated KV cache. Seed
                # safe defaults before probing so the except path has valid
                # state to publish.
                ctx_override = parse_ctx_override(extra_args)
                requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
                cache_override = parse_cache_override(extra_args)
                cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
                if ctx_override is not None and ctx_override > 0:
                    logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce")
                if cache_override is not None:
                    logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate")
                effective_ctx = requested_ctx if requested_ctx > 0 else (self._context_length or 0)
                max_available_ctx = self._context_length or effective_ctx
                gpus: list[tuple[int, int]] = []
                try:
                    model_size = self._get_gguf_size_bytes(model_path)
                    gpus = self._get_gpu_free_memory()

                    # Resolve effective context: 0 means let llama-server use
                    # the model's native length. Only expand to a known native
                    # length if metadata exists; else keep 0 as a sentinel.
                    if requested_ctx > 0:
                        effective_ctx = requested_ctx
                    elif self._context_length is not None:
                        effective_ctx = self._context_length
                    else:
                        effective_ctx = 0
                    original_ctx = effective_ctx
                    # Default UI ceiling to the native context length;
                    # GPU/VRAM-fit logic below may shrink it on limited HW.
                    max_available_ctx = self._context_length or effective_ctx

                    # Will MTP engage on this load? If so, auto-fit reserves
                    # extra VRAM for the draft model. Mirrors
                    # _build_speculative_flags' resolver: forced mtp / mtp+ngram
                    # always engage; auto only on an MTP GGUF >= 3B; ngram /
                    # ngram-simple / off never engage MTP.
                    _mtp_canonical = _canonicalize_spec_mode(speculative_type)
                    _mtp_effective = _mtp_canonical or "auto"
                    _mtp_size_for_fit = _extract_model_size_b(model_identifier)
                    _mtp_sub_3b_for_fit = _mtp_size_for_fit is not None and _mtp_size_for_fit < 3.0
                    _mtp_will_engage = bool(
                        not _extra_args_set_spec_type(extra_args)
                        and (
                            _mtp_effective in ("mtp", "mtp+ngram")
                            or (
                                _mtp_effective == "auto"
                                and (
                                    bool(self._nextn_predict_layers)
                                    or _is_mtp_model_name(model_identifier, model_path)
                                )
                                and not _mtp_sub_3b_for_fit
                            )
                        )
                    )

                    # Auto-cap context to fit GPU VRAM and select GPUs. Two
                    # policies by whether the user set n_ctx:
                    #   Explicit n_ctx: honor it. Try the full context with
                    #     _select_gpus (as many GPUs as needed); cap only if it
                    #     doesn't fit on any combination.
                    #   Auto n_ctx=0 (native): prefer fewer GPUs with reduced
                    #     context, since multi-GPU is slower.
                    gpu_indices, use_fit = None, True
                    explicit_ctx = requested_ctx > 0

                    if gpus and self._can_estimate_kv() and effective_ctx > 0:
                        # Largest hardware-aware cap from the native context
                        # across all usable GPU subsets (for UI bounds),
                        # independent of the requested context.
                        native_ctx_for_cap = self._context_length or effective_ctx
                        if native_ctx_for_cap > 0:
                            ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
                            best_cap = 0
                            for n_gpus in range(1, len(ranked_for_cap) + 1):
                                subset = ranked_for_cap[:n_gpus]
                                pool_mib = sum(free for _, free in subset)
                                capped = self._fit_context_to_vram(
                                    native_ctx_for_cap,
                                    pool_mib,
                                    model_size,
                                    cache_type_kv,
                                    n_parallel = n_parallel,
                                    mtp_engaged = _mtp_will_engage,
                                )
                                kv = self._estimate_kv_cache_bytes(
                                    capped, cache_type_kv, n_parallel = n_parallel
                                )
                                total_mib = (model_size + kv) / (1024 * 1024)
                                if total_mib <= pool_mib * 0.90:
                                    best_cap = max(best_cap, capped)
                            if best_cap > 0:
                                max_available_ctx = best_cap
                            else:
                                # Weights exceed 90% of every GPU subset, so no
                                # context fits. Anchor the UI "safe zone" at 4096
                                # so the slider warns above the fallback.
                                max_available_ctx = min(4096, native_ctx_for_cap)

                        if explicit_ctx:
                            # Honor the requested context verbatim. If it fits,
                            # pin GPUs and skip --fit; else ship -c <ctx> --fit
                            # on and let llama-server flex -ngl (CPU offload).
                            requested_total = model_size + self._estimate_kv_cache_bytes(
                                effective_ctx, cache_type_kv, n_parallel = n_parallel
                            )
                            gpu_indices, use_fit = self._select_gpus(requested_total, gpus)
                            # No silent shrink: effective_ctx stays == requested_ctx.
                        else:
                            # Auto context: prefer fewer GPUs, cap to fit. Same
                            # headroom threshold as _select_gpus (#5106).
                            ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
                            pin_fraction = self._GPU_PIN_VRAM_FRACTION
                            for n_gpus in range(1, len(ranked) + 1):
                                subset = ranked[:n_gpus]
                                pool_mib = sum(free for _, free in subset)
                                capped = self._fit_context_to_vram(
                                    effective_ctx,
                                    pool_mib,
                                    model_size,
                                    cache_type_kv,
                                    n_parallel = n_parallel,
                                    mtp_engaged = _mtp_will_engage,
                                )
                                kv = self._estimate_kv_cache_bytes(
                                    capped, cache_type_kv, n_parallel = n_parallel
                                )
                                total_mib = (model_size + kv) / (1024 * 1024)
                                if total_mib <= pool_mib * pin_fraction:
                                    effective_ctx = capped
                                    gpu_indices = sorted(idx for idx, _ in subset)
                                    use_fit = False
                                    break
                            else:
                                # Native ctx doesn't fit. Drop to 4096 and
                                # re-check before --fit on: a model overflowing
                                # at 131k may pin fine with a 4096 KV (#5106).
                                effective_ctx = min(4096, effective_ctx)
                                if effective_ctx > 0:
                                    for n_gpus in range(1, len(ranked) + 1):
                                        subset = ranked[:n_gpus]
                                        pool_mib = sum(free for _, free in subset)
                                        kv = self._estimate_kv_cache_bytes(
                                            effective_ctx,
                                            cache_type_kv,
                                            n_parallel = n_parallel,
                                        )
                                        total_mib = (model_size + kv) / (1024 * 1024)
                                        if total_mib <= pool_mib * pin_fraction:
                                            gpu_indices = sorted(idx for idx, _ in subset)
                                            use_fit = False
                                            break

                    elif gpus:
                        # Can't estimate KV -- file-size-only check; keep the
                        # ceiling at native context (already the default).
                        logger.debug(
                            "Falling back to file-size-only GPU selection",
                            model_size_gb = round(model_size / (1024**3), 2),
                        )
                        gpu_indices, use_fit = self._select_gpus(model_size, gpus)
                        if use_fit and not explicit_ctx:
                            # Weights don't fit on any subset; default UI to 4096
                            # so the slider isn't on an unusable native ctx.
                            effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096

                    if effective_ctx < original_ctx:
                        kv_est = self._estimate_kv_cache_bytes(
                            effective_ctx, cache_type_kv, n_parallel = n_parallel
                        )
                        logger.info(
                            f"Context auto-reduced: {original_ctx} -> {effective_ctx} "
                            f"(model: {model_size / (1024**3):.1f} GB, "
                            f"est. KV cache: {kv_est / (1024**3):.1f} GB)"
                        )

                    kv_cache_bytes = self._estimate_kv_cache_bytes(
                        effective_ctx, cache_type_kv, n_parallel = n_parallel
                    )
                    logger.info(
                        f"GGUF size: {model_size / (1024**3):.1f} GB, "
                        f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
                        f"context: {effective_ctx}, "
                        f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}"
                    )
                except Exception as e:
                    logger.warning(f"GPU selection failed ({e}), using --fit on")
                    gpu_indices, use_fit = None, True
                    effective_ctx = requested_ctx  # fall back to original

                launch_mmproj_path = self._resolve_launch_mmproj_path(
                    model_path = model_path,
                    mmproj_path = mmproj_path,
                )
                # Need both a resolved mmproj AND the config vision flag; a stray
                # mmproj passing the family-name heuristic must not flip a non-VLM
                # GGUF into vision mode.
                effective_is_vision = bool(launch_mmproj_path) and bool(is_vision)
                if is_vision and not effective_is_vision:
                    logger.warning(
                        "Vision-capable GGUF loaded without a usable mmproj; "
                        "image input will be disabled for this session"
                    )

                # Audio input straight from the mmproj (clip.has_audio_encoder),
                # independent of token names.
                self._mmproj_has_audio = False
                if launch_mmproj_path:
                    try:
                        from utils.models.gguf_metadata import (
                            read_mmproj_audio_capability,
                        )
                        self._mmproj_has_audio = bool(
                            read_mmproj_audio_capability(launch_mmproj_path)
                        )
                    except Exception as e:
                        logger.debug(f"mmproj audio-capability read failed: {e}")

                cmd = [
                    binary,
                    "-m",
                    model_path,
                    "--port",
                    str(self._port),
                    "-c",
                    str(effective_ctx) if effective_ctx > 0 else "0",
                    "--parallel",
                    str(n_parallel),
                    "--flash-attn",
                    "on",  # Force flash attention for speed
                    # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length".
                    "--no-context-shift",
                ]

                if use_fit:
                    cmd.extend(["--fit", "on"])
                elif gpu_indices is not None:
                    # Fits on selected GPU(s) -- offload all layers
                    cmd.extend(["-ngl", "-1"])

                # -1 = llama.cpp auto-detect (physical cores). Pass explicitly
                # so we don't inherit llama-server's internal default, which
                # has varied (hardware concurrency incl. hyperthreads on some
                # builds).
                cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)])

                # Enable Jinja chat template rendering
                cmd.extend(["--jinja"])

                # KV cache data type
                _valid_cache_types = {
                    "f16",
                    "bf16",
                    "q8_0",
                    "q4_0",
                    "q4_1",
                    "q5_0",
                    "q5_1",
                    "iq4_nl",
                    "f32",
                }
                if cache_type_kv and cache_type_kv in _valid_cache_types:
                    cmd.extend(
                        [
                            "--cache-type-k",
                            cache_type_kv,
                            "--cache-type-v",
                            cache_type_kv,
                        ]
                    )
                    self._cache_type_kv = cache_type_kv
                    logger.info(f"KV cache type: {cache_type_kv}")
                else:
                    self._cache_type_kv = None

                # Speculative decoding. See _build_speculative_flags for the
                # mode resolution, benchmarks, and llama.cpp references.
                spec_flags = self._build_speculative_flags(
                    speculative_type = speculative_type,
                    spec_draft_n_max = spec_draft_n_max,
                    extra_args = extra_args,
                    model_identifier = model_identifier,
                    model_path = model_path,
                    gpus = bool(gpus),
                    binary = binary,
                )
                cmd.extend(spec_flags)

                # Apply custom chat template override if provided.
                self._chat_template_override = chat_template_override
                if chat_template_override:
                    import tempfile

                    flags = detect_reasoning_flags(
                        chat_template_override,
                        self._model_identifier,
                        log_source = "GGUF chat template override",
                    )
                    self._supports_reasoning = flags["supports_reasoning"]
                    self._reasoning_style = flags["reasoning_style"]
                    self._reasoning_always_on = flags["reasoning_always_on"]
                    self._supports_preserve_thinking = flags["supports_preserve_thinking"]
                    self._supports_tools = flags["supports_tools"]

                    self._chat_template_file = tempfile.NamedTemporaryFile(
                        mode = "w",
                        suffix = ".jinja",
                        delete = False,
                        prefix = "unsloth_chat_template_",
                    )
                    self._chat_template_file.write(chat_template_override)
                    self._chat_template_file.close()
                    cmd.extend(["--chat-template-file", self._chat_template_file.name])
                    logger.info(f"Using custom chat template file: {self._chat_template_file.name}")

                # Default thinking mode for reasoning models. Qwen3.5/3.6 below
                # 9B disable thinking by default; 9B+ enable it. Always-on
                # templates ignore the kwarg, so skip.
                if self._supports_reasoning and not self._reasoning_always_on:
                    thinking_default = True
                    mid = (model_identifier or "").lower()
                    if "qwen3.5" in mid or "qwen3.6" in mid:
                        size_val = _extract_model_size_b(mid)
                        if size_val is not None and size_val < 9:
                            thinking_default = False
                    self._reasoning_default = thinking_default
                    reasoning_kw = self._reasoning_kwargs(thinking_default)
                    cmd.extend(
                        [
                            "--chat-template-kwargs",
                            json.dumps(reasoning_kw),
                        ]
                    )
                    logger.info(f"Reasoning model: {reasoning_kw} by default")

                if launch_mmproj_path and effective_is_vision:
                    cmd.extend(["--mmproj", launch_mmproj_path])
                    logger.info(f"Using mmproj for vision: {launch_mmproj_path}")

                # Option C: --api-key for direct client access when enabled
                import os as _os
                import secrets as _secrets

                if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1":
                    self._api_key = _secrets.token_urlsafe(32)
                    cmd.extend(["--api-key", self._api_key])
                    logger.info("llama-server started with --api-key for direct streaming")
                else:
                    self._api_key = None

                # User pass-through args go last so llama.cpp's last-wins parsing
                # lets the user override Studio's auto-set flags. Already
                # validated by the route via validate_extra_args().
                if extra_args:
                    cmd.extend(str(a) for a in extra_args)
                    logger.info(f"Appending user extra args to llama-server: {list(extra_args)}")

                _log_cmd = list(cmd)
                if "--api-key" in _log_cmd:
                    _ki = _log_cmd.index("--api-key") + 1
                    if _ki < len(_log_cmd):
                        _log_cmd[_ki] = "<redacted>"
                logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")

                # Library paths so llama-server finds its shared libs and CUDA DLLs.
                import os
                import sys

                env = child_env_without_native_path_secret()
                binary_dir = str(Path(binary).parent)

                # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
                # shared system RAM. setdefault so a user value wins.
                if self._amd_apu_wants_unified_memory():
                    env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
                    logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")

                if sys.platform == "win32":
                    # Ordering: see _build_windows_path_dirs. #5106.
                    path_dirs = self._build_windows_path_dirs(
                        binary_dir,
                        sys.prefix,
                        os.environ.get("CUDA_PATH", ""),
                    )
                    existing_path = env.get("PATH", "")
                    env["PATH"] = ";".join(path_dirs) + ";" + existing_path

                    # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile
                    # kernel files (rocblas/library/*.dat + *.hsaco); the DLL
                    # searches <binary_dir>/rocblas/library/ which doesn't exist
                    # -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH
                    # repoints that search at the ROCm install.
                    _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", ""))
                    if _hip_path:
                        _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library")
                        if os.path.isdir(_rocblas_lib):
                            env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
                else:
                    # Linux: LD_LIBRARY_PATH for shared libs next to the binary
                    # plus CUDA runtime libs (libcudart, libcublas, etc.)
                    import platform

                    lib_dirs = [binary_dir]
                    _arch = platform.machine()  # x86_64, aarch64, etc.

                    # Pip-installed nvidia CUDA runtime libs. The prebuilt
                    # binary links libcudart.so.13 / libcublas.so.13 which live
                    # here, not in /usr/local/cuda.
                    import glob as _glob

                    for _nv_pattern in [
                        os.path.join(
                            sys.prefix,
                            "lib",
                            "python*",
                            "site-packages",
                            "nvidia",
                            "cu*",
                            "lib",
                        ),
                        os.path.join(
                            sys.prefix,
                            "lib",
                            "python*",
                            "site-packages",
                            "nvidia",
                            "cudnn",
                            "lib",
                        ),
                        os.path.join(
                            sys.prefix,
                            "lib",
                            "python*",
                            "site-packages",
                            "nvidia",
                            "nvjitlink",
                            "lib",
                        ),
                    ]:
                        for _nv_dir in _glob.glob(_nv_pattern):
                            if os.path.isdir(_nv_dir):
                                lib_dirs.append(_nv_dir)

                    for cuda_lib in [
                        "/usr/local/cuda/lib64",
                        f"/usr/local/cuda/targets/{_arch}-linux/lib",
                        # Fallback CUDA compat paths (e.g. binary built with
                        # CUDA 12 where default /usr/local/cuda is CUDA 13+).
                        "/usr/local/cuda-12/lib64",
                        "/usr/local/cuda-12.8/lib64",
                        f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
                        f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
                    ]:
                        if os.path.isdir(cuda_lib):
                            lib_dirs.append(cuda_lib)
                    existing_ld = env.get("LD_LIBRARY_PATH", "")
                    new_ld = ":".join(lib_dirs)
                    env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld

                # Pin to selected GPU(s). On ROCm, narrowing only
                # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
                # HIP/ROCR set, so set those too.
                if gpu_indices is not None:
                    pinned = ",".join(str(i) for i in gpu_indices)
                    env["CUDA_VISIBLE_DEVICES"] = pinned
                    try:
                        import torch as _torch
                        if getattr(_torch.version, "hip", None) is not None:
                            env["HIP_VISIBLE_DEVICES"] = pinned
                            env["ROCR_VISIBLE_DEVICES"] = pinned
                    except Exception as e:
                        logger.debug("Failed to set ROCm visibility env vars for child: %s", e)

                # One-shot --fit off retry: recent llama.cpp runs a "fitting
                # params to device memory" step by default (--fit defaults to
                # 'on') even when -ngl is explicit. That step has aborted on
                # some ROCm hosts (ggml-cuda.cu ROCm error during worst-case
                # estimation, e.g. MTP + mmproj models on gfx1151). When
                # Studio's own VRAM math already placed the model
                # (use_fit=False), the step is redundant second-guessing --
                # retry once with --fit off before declaring the load failed.
                # Never retry when fit was requested (use_fit) or the caller
                # passed an explicit fit flag via extra args.
                _fit_retry_allowed = self._fit_off_retry_eligible(cmd, use_fit)
                for _spawn_attempt in (0, 1):
                    # Defensive kill: drop an orphan Popen a concurrent load may
                    # have stored before we overwrite the reference (#5161).
                    # Also reaps the crashed first attempt on the retry pass.
                    self._kill_process()

                    self._stdout_lines = []
                    # Tee llama-server output to a dedicated log file so a
                    # post-mortem has the full trail even when the parent only
                    # kept the last 50 lines. Path is under the studio home.
                    self._llama_log_fh = None
                    try:
                        log_dir = _swa_cache_path().parent / "logs" / "llama-server"
                        log_dir.mkdir(parents = True, exist_ok = True)
                        # Include the attempt index: the --fit off retry can
                        # respawn within the same epoch second, and reusing the
                        # name would truncate the crash log the retry warning
                        # just pointed the user at.
                        self._llama_log_path = log_dir / (
                            f"llama-{int(time.time())}-port-{self._port}"
                            f"-try{_spawn_attempt}.log"
                        )
                        self._llama_log_fh = open(
                            self._llama_log_path,
                            "w",
                            encoding = "utf-8",
                            buffering = 1,
                        )
                        logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
                    except OSError as e:
                        # Best-effort; never block the load on logging.
                        logger.debug(f"Could not open llama-server log file: {e}")
                        self._llama_log_path = None
                    self._process = subprocess.Popen(
                        cmd,
                        stdout = subprocess.PIPE,
                        stderr = subprocess.STDOUT,
                        text = True,
                        env = env,
                        **_windows_hidden_subprocess_kwargs(),
                    )

                    # Background thread to drain stdout (prevents pipe deadlock)
                    self._stdout_thread = threading.Thread(
                        target = self._drain_stdout, daemon = True, name = "llama-stdout"
                    )
                    self._stdout_thread.start()
                    if self._wait_for_health(timeout = 600.0):
                        break
                    _startup_crashed = (
                        self._process.poll() is not None and self._process.returncode != 0
                    )
                    if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed:
                        logger.warning(
                            "llama-server crashed during startup (exit code %s) "
                            "with the default memory-fit step enabled; Studio "
                            "already verified the model fits, retrying once "
                            "with --fit off. Crash log: %s",
                            self._process.returncode,
                            self._llama_log_path,
                        )
                        cmd = [*cmd, "--fit", "off"]
                        continue
                    self._kill_process()
                    raise RuntimeError(
                        self._classify_llama_start_failure(
                            "\n".join(self._stdout_lines[-50:]),
                            gguf_path,
                            model_identifier,
                        )
                    )

                # Store the resolved on-disk path, not the caller's kwarg: in
                # HF mode gguf_path is None and ``model_path`` is what
                # llama-server mmap's, which downstream consumers need.
                self._gguf_path = model_path
                self._hf_repo = hf_repo
                # For local GGUF files, extract variant from filename if absent
                if hf_variant:
                    self._hf_variant = hf_variant
                elif gguf_path:
                    try:
                        from utils.models.model_config import _extract_quant_label
                        self._hf_variant = _extract_quant_label(gguf_path)
                    except Exception:
                        self._hf_variant = None
                else:
                    self._hf_variant = None
                self._is_vision = effective_is_vision
                self._model_identifier = model_identifier

                # Store the effective (possibly capped) context separately; do
                # NOT overwrite _context_length (the native length for display).
                self._effective_context_length = (
                    effective_ctx if effective_ctx > 0 else self._context_length
                )
                self._max_context_length = (
                    max_available_ctx if max_available_ctx > 0 else self._effective_context_length
                )

                # Health was confirmed inside the spawn/retry loop above.
                self._healthy = True

                # Commit caller intent only after _healthy=True so a failed start
                # can't poison the next inheritance check. None keeps prior, []
                # clears, list sets. Source records hf_variant for the route's
                # same_source check.
                if extra_args is not None:
                    self._extra_args = list(extra_args)
                    self._extra_args_source = (model_identifier, hf_variant)
                self._requested_n_ctx = int(n_ctx)

                # Catch silent CPU fallback when GPU was intended (#5106).
                self._gpu_offload_active = self._classify_gpu_offload(
                    gpu_indices is not None or use_fit, gpus or []
                )
                if self._gpu_offload_active is False:
                    logger.warning(
                        "llama-server appears to have loaded the model entirely "
                        "on CPU even though Studio detected at least one GPU. "
                        "This usually means the prebuilt binary's GPU backend "
                        "failed to load -- on Windows, cudart64_X.dll / "
                        "cublas64_X.dll could not be resolved. Reinstall the "
                        "Studio llama.cpp prebuilt or install a matching CUDA "
                        "toolkit (issue unslothai/unsloth#5106).",
                    )

                logger.info(
                    f"llama-server ready on port {self._port} " f"for model '{model_identifier}'"
                )

            # Probe outside _lock (interruptible by /unload); init inside.
            self._is_audio = False
            self._audio_type = None
            self._audio_probed = False
            self._has_audio_input = False
            try:
                detected = self._detect_audio_type_strict()
                self._audio_probed = True
            except Exception as exc:
                logger.debug("Audio probe failed: %s", exc)
                detected = None
            if detected in ("snac", "bicodec", "dac"):
                with self._lock:
                    if not self._healthy:
                        return False
                    try:
                        self.init_audio_codec(detected)
                        self._is_audio = True
                        self._audio_type = detected
                    except Exception as exc:
                        # Surface as HTTP 500 (matches pre-PR contract).
                        logger.warning(
                            "Failed to init audio codec '%s': %s",
                            detected,
                            exc,
                        )
                        self._audio_probed = False
                        return False
            elif detected:
                # csm / whisper / audio_vlm: track type but keep _is_audio
                # False -- GGUF TTS routing only fires for snac/bicodec/dac.
                with self._lock:
                    if not self._healthy:
                        return False
                    self._audio_type = detected

            # Audio input = token probe (audio_vlm/whisper) OR mmproj encoder.
            from utils.models.model_config import is_audio_input_type

            self._has_audio_input = bool(is_audio_input_type(self._audio_type)) or bool(
                self._mmproj_has_audio
            )

            if not self._healthy:
                return False
            return True

    def _build_speculative_flags(
        self,
        *,
        speculative_type: Optional[str],
        spec_draft_n_max: Optional[int],
        extra_args: Optional[List[str]],
        model_identifier: str,
        model_path: Optional[str],
        gpus: bool,
        binary: Optional[str],
    ) -> List[str]:
        """Return the llama-server flag list for the requested spec mode.

        Side effects: sets ``self._speculative_type`` (resolved internal
        emit), ``self._requested_spec_mode`` (canonical UI mode for the
        status round-trip), and ``self._spec_draft_n_max`` (user override
        only; None when the platform default applies).

        Speculative decoding (n-gram self-speculation, zero VRAM):
        ngram-mod uses a ~16 MB shared hash pool, constant memory /
        complexity, variable draft lengths. Helps most when the model
        repeats existing text (code refactor, summarisation, reasoning);
        for low-repetition chat, overhead is ~5 ms.

        Benchmarks from upstream llama.cpp speculative-decoding PRs:
          Scenario                        | Without | With    | Speedup
          gpt-oss-120b code refactor      | 181 t/s | 446 t/s | 2.5x
          Qwen3-235B offloaded            |  12 t/s |  21 t/s | 1.8x
          gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
        Refs: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
              https://github.com/ggml-org/llama.cpp/pull/19164
              https://github.com/ggml-org/llama.cpp/pull/18471
              MTP guide: unsloth.ai/docs/models/qwen3.6#mtp-guide

        Sub-3B dense MTP regresses vs spec-off: the draft head's per-token
        cost exceeds the acceptance savings at this scale. Q4_K_XL clean
        bench (each prompt once after an unrelated warmup) on B200 + x86 CPU:
          0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x
          2B   GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x
          0.8B CPU: chained n=2   = 0.86x vs OFF; ngram-only = 1.19x
          2B   CPU: chained n=2   = 0.83x vs OFF; ngram-only = 1.01x
          4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
        Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
        diverse content); forced MTP variants engage anyway and just log a
        warning per the user's choice.
        """
        flags: List[str] = []
        # Reset; emit branches re-set on the resolved emission.
        self._spec_draft_n_max = None
        self._speculative_type = None

        # Canonical UI-facing requested mode (legacy values mapped via
        # _canonicalize_spec_mode).
        canonical_mode = _canonicalize_spec_mode(speculative_type)
        is_mtp_model = bool(self._nextn_predict_layers) or (
            _is_mtp_model_name(model_identifier, model_path)
        )
        user_owns_spec_type = _extra_args_set_spec_type(extra_args)
        _mtp_size_b = _extract_model_size_b(model_identifier)
        _mtp_too_small = _mtp_size_b is not None and _mtp_size_b < 3.0

        if user_owns_spec_type:
            # User --spec-type wins outright; suppress auto-emit to avoid a
            # duplicate spec block.
            self._requested_spec_mode = None
            return flags

        effective_mode = canonical_mode or "auto"
        self._requested_spec_mode = effective_mode

        def _resolved_draft_n_max() -> int:
            # User override wins; else platform default (the B200 / x86
            # clean-sweep sweet spot from PR #5582 is n=2 GPU, n=3 CPU;
            # past 3 regresses on essay-style low-acceptance prompts).
            if spec_draft_n_max is not None:
                n = int(spec_draft_n_max)
                self._spec_draft_n_max = n
                return n
            return 2 if gpus else 3

        def _emit_mtp(*, chain_ngram: bool) -> bool:
            """Append --spec-type mtp[/draft-mtp][,ngram-mod] + n-max."""
            caps = self.probe_server_capabilities(binary)
            mtp_token = caps.get("mtp_token") if caps else None
            if not mtp_token:
                logger.warning(
                    "Requested MTP speculative decoding but "
                    "llama-server lacks --spec-type mtp/draft-mtp; "
                    "run `unsloth studio update`. Loading without "
                    "speculative decoding."
                )
                return False
            draft_n_max = _resolved_draft_n_max()
            n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
            if chain_ngram:
                ngram_knobs = _build_ngram_mod_flags(caps)
                if ngram_knobs:
                    spec_value = f"ngram-mod,{mtp_token}"
                else:
                    logger.warning(
                        "llama-server lacks ngram-mod tuning "
                        "flags; loading MTP only (no ngram chain)"
                    )
                    spec_value = mtp_token
                flags.extend(
                    [
                        "--spec-type",
                        spec_value,
                        n_max_flag,
                        str(draft_n_max),
                    ]
                )
                flags.extend(ngram_knobs)
            else:
                flags.extend(
                    [
                        "--spec-type",
                        mtp_token,
                        n_max_flag,
                        str(draft_n_max),
                    ]
                )
            self._speculative_type = "draft-mtp"
            chain_label = "chained ngram-mod" if chain_ngram else "MTP-only"
            logger.info(f"Spec decoding: {mtp_token} ({chain_label})")
            return True

        def _emit_ngram_mod() -> bool:
            """Append --spec-type ngram-mod + flag-set knobs."""
            ngram_caps = self.probe_server_capabilities(binary)
            ngram_knobs = _build_ngram_mod_flags(ngram_caps)
            flags.extend(["--spec-type", "ngram-mod"])
            if not ngram_knobs:
                logger.warning(
                    "llama-server lacks ngram-mod tuning "
                    "flags; loading without --spec-ngram-mod-* knobs"
                )
            flags.extend(ngram_knobs)
            self._speculative_type = "ngram-mod"
            logger.info("Spec decoding: ngram-mod")
            return True

        if effective_mode == "off":
            return flags  # nothing to emit
        if effective_mode == "ngram-simple":
            flags.extend(["--spec-type", "ngram-simple"])
            self._speculative_type = "ngram-simple"
            return flags
        if effective_mode == "ngram":
            _emit_ngram_mod()
            return flags
        if effective_mode == "mtp":
            if _mtp_too_small:
                logger.warning(
                    f"Forcing MTP on a {_mtp_size_b:.1f}B model; "
                    "the bench shows draft-mtp regresses below 3B. "
                    "Engaging anyway (user override)."
                )
            elif not is_mtp_model:
                logger.warning(
                    "Forcing MTP on a non-MTP GGUF; llama-server may "
                    "fall back to spec-off if no nextn head is present. "
                    "Engaging anyway (user override)."
                )
            _emit_mtp(chain_ngram = False)
            return flags
        if effective_mode == "mtp+ngram":
            if _mtp_too_small:
                logger.warning(
                    f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; "
                    "the bench shows the chain regresses below 3B. "
                    "Engaging anyway (user override)."
                )
            elif not is_mtp_model:
                logger.warning(
                    "Forcing MTP+Ngram on a non-MTP GGUF; llama-server "
                    "may fall back to ngram-only if no nextn head is "
                    "present. Engaging anyway (user override)."
                )
            _emit_mtp(chain_ngram = True)
            return flags

        # effective_mode == "auto": the promotion path. llama.cpp #22673:
        # MTP is compatible with mmproj, so there's no vision gate.
        if is_mtp_model and not _mtp_too_small:
            # GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP.
            _emit_mtp(chain_ngram = not gpus)
        elif is_mtp_model and _mtp_too_small:
            # Sub-3B fallback: drop the MTP draft head, keep ngram-mod when
            # the binary supports it.
            _small_caps = self.probe_server_capabilities(binary)
            if _small_caps.get("supports_ngram_mod"):
                logger.info(
                    f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
                    "is below the 3B speedup threshold; using ngram-mod "
                    "only (zero-VRAM, no draft head). Override via "
                    "--spec-type or the Studio Speculative Decoding "
                    "dropdown."
                )
                _emit_ngram_mod()
            else:
                logger.info(
                    f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
                    "is below the 3B speedup threshold and the bundled "
                    "llama-server does not advertise ngram-mod; "
                    "auto-disabling speculative decoding."
                )
        else:
            # Non-MTP model: let llama-server choose its default strategy.
            flags.append("--spec-default")
            self._speculative_type = "default"
        return flags

    def _already_in_target_state(
        self,
        *,
        model_identifier: str,
        hf_variant: Optional[str],
        n_ctx: int,
        cache_type_kv: Optional[str],
        speculative_type: Optional[str],
        chat_template_override: Optional[str],
        extra_args: Optional[List[str]],
        is_vision: bool,
        gguf_path: Optional[str] = None,
        spec_draft_n_max: Optional[int] = None,
    ) -> bool:
        """True iff the live server already satisfies these load kwargs.

        Mirrors ``routes/inference.py:_request_matches_loaded_settings`` but
        compares raw kwargs so ``load_model`` can short-circuit a duplicate
        /load that raced past the route-level check (#5401).
        """
        if not self.is_loaded:
            return False
        if (self._model_identifier or "").lower() != (model_identifier or "").lower():
            return False
        # Direct-file loads pass hf_variant=None while the backend stores an
        # extracted filename label; compare paths to keep the guard symmetric.
        if gguf_path is not None and self._gguf_path:
            try:
                if Path(self._gguf_path).resolve() != Path(gguf_path).resolve():
                    return False
            except OSError:
                return False
        elif (self._hf_variant or "").lower() != (hf_variant or "").lower():
            return False
        if self._requested_n_ctx != int(n_ctx):
            return False

        def _norm(value):
            if value is None:
                return None
            if isinstance(value, str):
                stripped = value.strip().lower()
                return stripped or None
            return value

        if _norm(self._cache_type_kv) != _norm(cache_type_kv):
            return False

        # Compare on the canonical requested mode. With --spec-type in
        # extra_args the backend stores None; mirror that here.
        if _extra_args_set_spec_type(extra_args):
            req_mode = None
        else:
            req_mode = _canonicalize_spec_mode(speculative_type) or "auto"
        backend_mode = self._requested_spec_mode
        if req_mode != backend_mode:
            return False

        # spec_draft_n_max only matters when an MTP variant is engaged. Compare
        # on the resolved spec so an Auto request promoted to draft-mtp still
        # bounces a reload when n_max changes.
        if (
            self._speculative_type == "draft-mtp"
            and spec_draft_n_max is not None
            and int(spec_draft_n_max) != (self._spec_draft_n_max or 0)
        ):
            return False

        if (self._chat_template_override or None) != (chat_template_override or None):
            return False

        # extra_args=None means "no opinion" (inherit handled at the route
        # layer); only an explicit list forces equality.
        if extra_args is not None:
            current = list(self._extra_args) if self._extra_args is not None else []
            if list(extra_args) != current:
                return False
        return True

    def _classify_gpu_offload(
        self, expected_gpu: bool, detected_gpus: list[tuple[int, int]]
    ) -> Optional[bool]:
        """True if a GPU model buffer was allocated, False if only CPU
        buffers landed despite GPU intent, None when there's no signal (no
        GPU detected, no buffer-size lines, etc.)."""
        if not detected_gpus or not expected_gpu:
            return None
        # llama-server logs one "model buffer size = N MiB" line per backend
        # buffer; CUDA/ROCm/Metal/Vulkan/OpenCL/SYCL are GPU, CPU* are not.
        gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL")
        saw_buffer_line = False
        saw_gpu_buffer = False
        for line in self._stdout_lines:
            if "model buffer size" not in line:
                continue
            saw_buffer_line = True
            if any(marker in line for marker in gpu_markers):
                saw_gpu_buffer = True
                break
        if not saw_buffer_line:
            return None
        return saw_gpu_buffer

    def unload_model(self) -> bool:
        """Terminate the subprocess and cancel any in-flight download."""
        self._cancel_event.set()
        with self._lock:
            self._kill_process()
            logger.info(f"Unloaded GGUF model: {self._model_identifier}")
            self._model_identifier = None
            self._gguf_path = None
            self._hf_repo = None
            self._hf_variant = None
            self._is_vision = False
            self._is_audio = False
            self._audio_type = None
            self._audio_probed = False
            self._has_audio_input = False
            self._mmproj_has_audio = False
            self._port = None
            self._healthy = False
            self._context_length = None
            self._effective_context_length = None
            self._max_context_length = None
            self._chat_template = None
            self._chat_template_override = None
            self._supports_reasoning = False
            self._reasoning_always_on = False
            self._reasoning_style = "enable_thinking"
            self._reasoning_default = True
            self._supports_preserve_thinking = False
            self._supports_tools = False
            self._cache_type_kv = None
            self._speculative_type = None
            self._requested_spec_mode = None
            self._spec_draft_n_max = None
            self._n_layers = None
            self._n_kv_heads = None
            self._n_kv_heads_by_layer = None
            self._n_heads = None
            self._embedding_length = None
            self._kv_key_length = None
            self._kv_value_length = None
            self._sliding_window = None
            self._sliding_window_pattern = None
            self._full_attention_interval = None
            self._kv_lora_rank = None
            self._key_length_mla = None
            self._kv_key_length_swa = None
            self._kv_value_length_swa = None
            self._ssm_inner_size = None
            self._ssm_state_size = None
            self._shared_kv_layers = None
            self._nextn_predict_layers = None
            # Clean up temp chat template file.
            if hasattr(self, "_chat_template_file") and self._chat_template_file:
                try:
                    import os
                    os.unlink(self._chat_template_file.name)
                except Exception:
                    pass
                self._chat_template_file = None
            # Free audio codec GPU memory.
            if LlamaCppBackend._codec_mgr is not None:
                LlamaCppBackend._codec_mgr.unload()
                LlamaCppBackend._codec_mgr = None
                import torch

                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            return True

    def _kill_process(self):
        """Terminate the subprocess if running."""
        if self._process is None:
            return
        try:
            self._process.terminate()
            self._process.wait(timeout = 5)
        except subprocess.TimeoutExpired:
            logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
            self._process.kill()
            self._process.wait(timeout = 5)
        except Exception as e:
            logger.warning(f"Error killing llama-server process: {e}")
        finally:
            self._process = None
            # Clear healthy so a /load during the replacement's warm-up can't
            # short-circuit against the previous server's health (#5401).
            self._healthy = False
            # Drives _wait_for_vram_settle in the next load_model; set in finally
            # so both in-process and frontend Apply paths record the kill.
            self._last_kill_monotonic = time.monotonic()
            if self._stdout_thread is not None:
                self._stdout_thread.join(timeout = 2)
                self._stdout_thread = None
            fh = getattr(self, "_llama_log_fh", None)
            if fh is not None:
                try:
                    fh.close()
                except Exception:
                    pass
                self._llama_log_fh = None

    @staticmethod
    def _kill_orphaned_servers():
        """Kill orphaned llama-server processes started by studio.

        Only kills processes whose resolved binary lives under a known
        Studio install dir (or matches an exact env-var override), to avoid
        terminating unrelated llama-server instances. Mirrors every location
        _find_llama_server_binary() can return, so orphans from any
        supported install path are cleaned up.

        Uses psutil for cross-platform support (Linux, macOS, Windows);
        falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
        absent.
        """
        try:
            # -- Build the ownership allowlist --------------------------------
            # exact_binaries -- env var overrides (exact path match).
            # install_roots  -- Studio-owned dir trees (binary must be under one).
            install_roots: list[Path] = []

            # Env-mode custom root (mirrors _find_llama_server_binary).
            _is_custom_root = False
            try:
                from utils.paths.storage_roots import studio_root as _sr  # noqa: WPS433

                _resolved_sr = _sr()
                _legacy_studio = Path.home() / ".unsloth" / "studio"
                try:
                    _is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve()
                except (OSError, ValueError):
                    _is_custom_root = _resolved_sr != _legacy_studio
                if _is_custom_root:
                    install_roots.append(_resolved_sr / "llama.cpp")
            except (ImportError, OSError, ValueError):
                pass

            # Primary install dir (default mode only). Env-mode skips this so a
            # custom-root Studio can't kill a default-install Studio's server.
            if not _is_custom_root:
                install_roots.append(Path.home() / ".unsloth" / "llama.cpp")

            # Legacy in-tree build dirs (older setup.sh)
            project_root = Path(__file__).resolve().parents[4]
            install_roots.append(project_root / "llama.cpp")

            # Legacy: extracted binary
            install_roots.append(project_root / "bin")

            # UNSLOTH_LLAMA_CPP_PATH env var (custom install dir)
            custom_dir = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
            if custom_dir:
                install_roots.append(Path(custom_dir))

            # LLAMA_SERVER_PATH env var (exact binary path)
            exact_binaries: list[Path] = []
            env_binary = os.environ.get("LLAMA_SERVER_PATH")
            if env_binary:
                try:
                    exact_binaries.append(Path(env_binary).resolve())
                except OSError:
                    pass

            # Resolve all roots so is_relative_to works reliably.
            resolved_roots: list[Path] = []
            for root in install_roots:
                try:
                    resolved_roots.append(root.resolve())
                except OSError:
                    pass

            my_pid = os.getpid()

            # -- Enumerate processes -------------------------------------------
            # Prefer psutil (cross-platform); fall back to pgrep + /proc on
            # Linux when psutil is absent.
            try:
                import psutil
                has_psutil = True
            except ImportError:
                has_psutil = False

            if has_psutil:
                for proc in psutil.process_iter(["pid", "name", "exe"]):
                    try:
                        if proc.info["pid"] == my_pid:
                            continue

                        name = proc.info.get("name") or ""
                        if not name.lower().startswith("llama-server"):
                            continue

                        exe = proc.info.get("exe")
                        if not exe:
                            continue

                        exe_path = Path(exe).resolve()

                        # Ownership: exact match OR binary under a known root.
                        is_ours = exe_path in exact_binaries or any(
                            exe_path.is_relative_to(root) for root in resolved_roots
                        )
                        if not is_ours:
                            continue

                        proc.kill()
                        logger.info(
                            f"Killed orphaned llama-server process " f"(pid={proc.info['pid']})"
                        )
                    except (
                        psutil.NoSuchProcess,
                        psutil.AccessDenied,
                        psutil.ZombieProcess,
                    ):
                        pass
            else:
                # -- Fallback: pgrep + /proc/<pid>/exe (Linux only) -----------
                if sys.platform != "linux":
                    return
                result = subprocess.run(
                    ["pgrep", "-a", "-f", "llama-server"],
                    capture_output = True,
                    text = True,
                    timeout = 5,
                    env = child_env_without_native_path_secret(),
                )
                if result.returncode != 0:
                    return

                for line in result.stdout.strip().splitlines():
                    parts = line.strip().split(None, 1)
                    if len(parts) < 2:
                        continue
                    pid = int(parts[0])
                    if pid == my_pid:
                        continue

                    # /proc/<pid>/exe symlinks the real binary, avoiding
                    # cmdline-parsing ambiguities; fall back to the first
                    # cmdline token when /proc is unavailable.
                    proc_exe = Path(f"/proc/{pid}/exe")
                    try:
                        binary = proc_exe.resolve(strict = True)
                    except (OSError, ValueError):
                        cmdline = parts[1]
                        token = cmdline.split()[0] if cmdline.strip() else ""
                        if not token:
                            continue
                        binary = Path(token).resolve(strict = False)

                    owned = binary in exact_binaries or any(
                        binary.is_relative_to(root) for root in resolved_roots
                    )
                    if not owned:
                        continue

                    try:
                        os.kill(pid, signal.SIGKILL)
                        logger.info(f"Killed orphaned llama-server process (pid={pid})")
                    except ProcessLookupError:
                        pass
                    except PermissionError:
                        pass
        except Exception:
            logger.warning("Error during orphan server cleanup", exc_info = True)

    def _cleanup(self):
        """atexit handler to ensure llama-server is terminated."""
        self._kill_process()

    @staticmethod
    def _fit_off_retry_eligible(cmd: "list[str]", use_fit: bool) -> bool:
        """Whether a llama-server startup crash may be retried with --fit off.

        Only when Studio's own VRAM math placed the model (use_fit=False)
        and nothing on the command line set the fit mode explicitly
        (-fit / --fit, space- or equals-form). --fit-ctx / --fit-target /
        -fitc / -fitt tune the fit step but do not select the mode, so
        they do not block the retry.
        """
        if use_fit:
            return False
        for a in cmd:
            if a in ("-fit", "--fit") or a.startswith(("-fit=", "--fit=")):
                return False
        return True

    def _wait_for_health(
        self,
        timeout: float = 120.0,
        interval: float = 0.5,
    ) -> bool:
        """Poll llama-server's /health until 200; also detect early exit/crash."""
        deadline = time.monotonic() + timeout
        url = f"http://127.0.0.1:{self._port}/health"

        while time.monotonic() < deadline:
            # Process crashed?
            if self._process.poll() is not None:
                # Let the drain thread collect final output.
                if self._stdout_thread is not None:
                    self._stdout_thread.join(timeout = 2)
                output = "\n".join(self._stdout_lines[-50:])
                # Keep the TAIL: crash details (abort reason, ROCm/CUDA error
                # text) print last, after the long startup banner. Head
                # truncation has cut off exactly the diagnostic line before.
                _log_hint = (
                    f" Full log: {self._llama_log_path}"
                    if getattr(self, "_llama_log_path", None)
                    else ""
                )
                logger.error(
                    f"llama-server exited with code {self._process.returncode}. "
                    f"Output (tail): {output[-2000:]}{_log_hint}"
                )
                return False

            try:
                resp = httpx.get(url, timeout = 2.0)
                if resp.status_code == 200:
                    return True
            except (
                httpx.ConnectError,
                httpx.TimeoutException,
                # ReadError covers TCP RST mid-read while still binding the port
                # (Windows: WinError 10054); the crash branch catches real exits.
                httpx.ReadError,
                httpx.RemoteProtocolError,
                httpx.WriteError,
            ):
                pass

            time.sleep(interval)

        logger.error(f"llama-server health check timed out after {timeout}s")
        return False

    # ── Message building (OpenAI format) ──────────────────────────

    @staticmethod
    def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]:
        """Thin wrapper around the shared parser in tool_call_parser
        so safetensors and llama_cpp pick up the same fixes."""
        return _shared_parse_tool_calls_from_text(
            content,
            allow_incomplete = allow_incomplete,
        )

    @staticmethod
    def _build_openai_messages(messages: list[dict], image_b64: Optional[str] = None) -> list[dict]:
        """Build OpenAI-format messages, optionally injecting an image_url part
        into the last user message for vision models. As-is if no image."""
        if not image_b64:
            return messages

        # Convert the last user message to multimodal content parts
        result = [msg.copy() for msg in messages]
        last_user_idx = None
        for i, msg in enumerate(result):
            if msg["role"] == "user":
                last_user_idx = i

        if last_user_idx is not None:
            text_content = result[last_user_idx].get("content", "")
            result[last_user_idx]["content"] = [
                {"type": "text", "text": text_content},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{image_b64}",
                    },
                },
            ]

        return result

    # ── Generation (proxy to llama-server) ────────────────────────

    @staticmethod
    def _iter_text_cancellable(
        response: "httpx.Response", cancel_event: Optional[threading.Event] = None
    ) -> Generator[str, None, None]:
        """Iterate an httpx streaming response with cancel support.

        Checks cancel_event between chunks and on ReadTimeout; the
        _stream_with_retry watcher also closes the response on cancel.
        """
        text_iter = response.iter_text()
        while True:
            if cancel_event is not None and cancel_event.is_set():
                response.close()
                return
            try:
                chunk = next(text_iter)
                yield chunk
            except StopIteration:
                return
            except httpx.ReadTimeout:
                # No data within the timeout window -- loop back and re-check
                # cancel_event.
                continue

    @staticmethod
    @contextlib.contextmanager
    def _stream_with_retry(
        client: "httpx.Client",
        url: str,
        payload: dict,
        cancel_event: Optional[threading.Event] = None,
        headers: Optional[dict] = None,
    ):
        """Open an httpx streaming POST with cancel support.

        Sends once with a long read timeout (120 s) so prefill finishes without
        a retry storm (the old 0.5 s timeout caused duplicate POSTs every half
        second). A watcher thread cancels by closing the response. httpx can't
        interrupt a blocked read before the response exists, so cancel during
        the header wait (1-5 s prefill) is deferred until headers arrive.
        """
        if cancel_event is not None and cancel_event.is_set():
            raise GeneratorExit

        # Background watcher: close the response if cancel is requested.
        # Only effective after response headers arrive (httpx limitation).
        _cancel_closed = threading.Event()
        _response_ref: list = [None]

        def _cancel_watcher():
            while not _cancel_closed.is_set():
                if cancel_event.wait(timeout = 0.3):
                    # Cancel requested. Poll until the response object exists
                    # so we can close it, or until the main thread finishes
                    # (_cancel_closed set in finally).
                    while not _cancel_closed.is_set():
                        r = _response_ref[0]
                        if r is not None:
                            try:
                                r.close()
                                return
                            except Exception as e:
                                logger.debug(f"Error closing response in cancel watcher: {e}")
                        # Response not created yet -- wait briefly and retry
                        _cancel_closed.wait(timeout = 0.1)
                    return

        watcher = None
        if cancel_event is not None:
            watcher = threading.Thread(target = _cancel_watcher, daemon = True, name = "prefill-cancel")
            watcher.start()

        try:
            # Long read timeout so prefill can finish without a retry storm.
            # Cancel during prefill and streaming is handled by the watcher
            # thread closing the response, unblocking any httpx read.
            prefill_timeout = httpx.Timeout(
                connect = 30,
                read = 120.0,
                write = 10,
                pool = 10,
            )
            with client.stream(
                "POST",
                url,
                json = payload,
                timeout = prefill_timeout,
                headers = headers,
            ) as response:
                _response_ref[0] = response
                if cancel_event is not None and cancel_event.is_set():
                    raise GeneratorExit
                yield response
                return
        except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError):
            # Response was closed by the cancel watcher
            if cancel_event is not None and cancel_event.is_set():
                raise GeneratorExit
            raise
        finally:
            _cancel_closed.set()

    def generate_chat_completion(
        self,
        messages: list[dict],
        image_b64: Optional[str] = None,
        temperature: float = 0.6,
        top_p: float = 0.95,
        top_k: int = 20,
        min_p: float = 0.01,
        max_tokens: Optional[int] = None,
        repetition_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        stop: Optional[list[str]] = None,
        cancel_event: Optional[threading.Event] = None,
        enable_thinking: Optional[bool] = None,
        reasoning_effort: Optional[str] = None,
        preserve_thinking: Optional[bool] = None,
        seed: Optional[int] = None,
    ) -> Generator[str | dict, None, None]:
        """
        Send a chat completion to llama-server and stream tokens back.

        Uses /v1/chat/completions -- llama-server applies the chat template
        and handles vision (multimodal image_url parts) natively.

        Yields cumulative text (matching InferenceBackend's convention).
        """
        if not self.is_loaded:
            raise RuntimeError("llama-server is not loaded")

        openai_messages = self._build_openai_messages(messages, image_b64)

        payload = {
            "messages": openai_messages,
            "stream": True,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k if top_k >= 0 else 0,
            "min_p": min_p,
            "repeat_penalty": repetition_penalty,
            "presence_penalty": presence_penalty,
        }
        # Per-request enable_thinking / reasoning_effort / preserve_thinking
        _reasoning_kw = self._request_reasoning_kwargs(
            enable_thinking, reasoning_effort, preserve_thinking
        )
        if _reasoning_kw is not None:
            payload["chat_template_kwargs"] = _reasoning_kw
        # Cap to the effective context length when known, else the floor.
        # The wall-clock backstop below stops a stuck model regardless.
        payload["max_tokens"] = (
            max_tokens
            if max_tokens is not None
            else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
        )
        payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
        if stop:
            payload["stop"] = stop
        if seed is not None:
            payload["seed"] = seed
        payload["stream_options"] = {"include_usage": True}

        url = f"{self.base_url}/v1/chat/completions"
        cumulative = ""
        in_thinking = False
        _stream_done = False
        _metadata_usage = None
        _metadata_timings = None
        _metadata_finish_reason = None

        try:
            # _stream_with_retry uses a 120 s read timeout so prefill can
            # finish. Cancel during streaming is handled by the watcher
            # thread (closes the response on cancel_event).
            stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
            _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
            with httpx.Client(
                timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
            ) as client:
                with self._stream_with_retry(
                    client,
                    url,
                    payload,
                    cancel_event,
                    headers = _auth_headers,
                ) as response:
                    if response.status_code != 200:
                        error_body = response.read().decode()
                        raise RuntimeError(
                            f"llama-server returned {response.status_code}: {error_body}"
                        )

                    buffer = ""
                    has_content_tokens = False
                    reasoning_text = ""
                    for raw_chunk in self._iter_text_cancellable(response, cancel_event):
                        buffer += raw_chunk
                        while "\n" in buffer:
                            line, buffer = buffer.split("\n", 1)
                            line = line.strip()

                            if not line:
                                continue
                            if line == "data: [DONE]":
                                if in_thinking:
                                    if has_content_tokens:
                                        # Real thinking + content: close the tag
                                        cumulative += "</think>"
                                        yield cumulative
                                    else:
                                        # Only reasoning_content, no content:
                                        # model put its whole reply in reasoning
                                        # (e.g. Qwen3 always-think). Show it as
                                        # the main response, not a thinking block.
                                        cumulative = reasoning_text
                                        yield cumulative
                                _stream_done = True
                                break  # exit inner while
                            if not line.startswith("data: "):
                                continue

                            try:
                                data = json.loads(line[6:])
                                # Capture server timings/usage from final chunks.
                                _chunk_timings = data.get("timings")
                                if _chunk_timings:
                                    _metadata_timings = _chunk_timings
                                _chunk_usage = data.get("usage")
                                if _chunk_usage:
                                    _metadata_usage = _chunk_usage
                                choices = data.get("choices", [])
                                if choices:
                                    delta = choices[0].get("delta", {})
                                    _fr = choices[0].get("finish_reason")
                                    if _fr:
                                        _metadata_finish_reason = _fr

                                    # Reasoning/thinking tokens: llama-server
                                    # sends these as "reasoning_content"; wrap
                                    # in <think> tags for the frontend parser.
                                    reasoning = delta.get("reasoning_content", "")
                                    if reasoning:
                                        reasoning_text += reasoning
                                        if not in_thinking:
                                            cumulative += "<think>"
                                            in_thinking = True
                                        cumulative += reasoning
                                        yield cumulative

                                    token = delta.get("content", "")
                                    if token:
                                        has_content_tokens = True
                                        if in_thinking:
                                            cumulative += "</think>"
                                            in_thinking = False
                                        cumulative += token
                                        yield cumulative
                            except json.JSONDecodeError:
                                logger.debug(f"Skipping malformed SSE line: {line[:100]}")
                        if _stream_done:
                            break  # exit outer for
                    if _metadata_usage or _metadata_timings or _metadata_finish_reason:
                        _metadata_usage = _backfill_usage_from_timings(
                            _metadata_usage, _metadata_timings
                        )
                        yield {
                            "type": "metadata",
                            # Never None: a finish-only metadata event (no usage,
                            # no timings) would otherwise crash consumers that do
                            # usage.get(...) on the non-streaming paths.
                            "usage": _metadata_usage or {},
                            "timings": _metadata_timings,
                            "finish_reason": _metadata_finish_reason,
                        }

        except httpx.ConnectError:
            raise RuntimeError("Lost connection to llama-server")
        except Exception as e:
            if cancel_event is not None and cancel_event.is_set():
                return
            raise

    # ── Tool-calling agentic loop ──────────────────────────────

    def generate_chat_completion_with_tools(
        self,
        messages: list[dict],
        tools: list[dict],
        temperature: float = 0.6,
        top_p: float = 0.95,
        top_k: int = 20,
        min_p: float = 0.01,
        max_tokens: Optional[int] = None,
        repetition_penalty: float = 1.0,
        presence_penalty: float = 0.0,
        stop: Optional[list[str]] = None,
        cancel_event: Optional[threading.Event] = None,
        enable_thinking: Optional[bool] = None,
        reasoning_effort: Optional[str] = None,
        preserve_thinking: Optional[bool] = None,
        max_tool_iterations: int = 25,
        auto_heal_tool_calls: bool = True,
        tool_call_timeout: int = 300,
        session_id: Optional[str] = None,
        rag_scope: Optional[dict] = None,
        seed: Optional[int] = None,
        disable_parallel_tool_use: bool = False,
    ) -> Generator[dict, None, None]:
        """
        Agentic loop: let the model call tools, execute them, and continue.

        Yields dicts:
          {"type": "status", "text": "Searching: ..."/"Reading: ..."}   -- tool status updates
          {"type": "content", "text": "token"}            -- streamed content tokens (cumulative)
          {"type": "reasoning", "text": "token"}          -- streamed reasoning tokens (cumulative)
        """
        from core.inference.tools import build_rag_autoinject, execute_tool

        if not self.is_loaded:
            raise RuntimeError("llama-server is not loaded")

        conversation = list(messages)

        # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits
        # the same tool card + citations a real call would.
        _auto = build_rag_autoinject(conversation, rag_scope)
        if _auto:
            for _ev in _auto["events"]:
                yield _ev
            conversation.extend(_auto["messages"])

        url = f"{self.base_url}/v1/chat/completions"
        _accumulated_completion_tokens = 0
        _accumulated_predicted_ms = 0.0
        _accumulated_predicted_n = 0

        def _strip_tool_markup(
            text: str,
            *,
            final: bool = False,
            force: bool = False,
        ) -> str:
            if not (auto_heal_tool_calls or force):
                return text
            return strip_tool_call_markup(text, final = final)

        def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
            if not (auto_heal_tool_calls or force):
                return text
            for pat in _TOOL_ALL_PATS:
                text = pat.sub("", text)
            return text

        tool_controller = ToolLoopController(
            tools = tools,
            auto_heal_tool_calls = auto_heal_tool_calls,
        )

        def _tool_succeeded(tool_name: str) -> bool:
            key_prefix = f"{tool_name}:"
            return any(
                record.executed and not record.is_error and record.key.startswith(key_prefix)
                for record in tool_controller.history
            )

        _MAX_BUFFER_CHARS = 32
        _append_budget_exhausted_nudge = True
        # RAG: cap knowledge-base searches per assistant turn. The controller is
        # tool-agnostic, so this gate stays in the loop.
        _kb_search_count = 0

        # ── Re-prompt on plan-without-action ─────────────────
        # When the model describes what it intends to do (forward-looking
        # language) without calling a tool, re-prompt once. Only triggers on
        # responses signaling intent/planning -- a direct answer like "4" or
        # "Hello!" won't match. Pattern compiled at module level
        # (_INTENT_SIGNAL).
        _reprompt_count = 0
        _forced_tool_call_pending = False

        # Reserve extra iterations for re-prompts so they don't consume the
        # caller's tool-call budget; only when tool iterations are allowed.
        _extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
        for iteration in range(max_tool_iterations + _extra):
            if cancel_event is not None and cancel_event.is_set():
                return

            active_tools = tool_controller.active_tools()
            if not active_tools:
                _append_budget_exhausted_nudge = False
                break
            _tool_xml_signals = TOOL_XML_SIGNALS if active_tools else ()

            # Build payload -- stream: True so we detect tool signals
            # in the first 1-2 chunks without a non-streaming penalty.
            payload = {
                "messages": conversation,
                "stream": True,
                "stream_options": {"include_usage": True},
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k if top_k >= 0 else 0,
                "min_p": min_p,
                "repeat_penalty": repetition_penalty,
                "presence_penalty": presence_penalty,
                "tools": active_tools,
                "tool_choice": "auto",
            }
            _reasoning_kw = self._request_reasoning_kwargs(
                enable_thinking, reasoning_effort, preserve_thinking
            )
            if _reasoning_kw is not None:
                payload["chat_template_kwargs"] = _reasoning_kw
            payload["max_tokens"] = (
                max_tokens
                if max_tokens is not None
                else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
            )
            payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
            if stop:
                payload["stop"] = stop
            if seed is not None:
                payload["seed"] = seed

            try:
                _auth_headers = (
                    {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
                )

                # ── Speculative buffer state machine ──────────────────
                # BUFFERING: accumulate content, check for tool signals
                # STREAMING: no tool detected, yield tokens to caller
                # DRAINING:  tool signal found, silently consume rest
                _S_BUFFERING = 0
                _S_STREAMING = 1
                _S_DRAINING = 2

                detect_state = _S_BUFFERING
                content_buffer = ""  # Raw content held during BUFFERING
                content_accum = ""  # All content tokens (for tool parsing)
                reasoning_accum = ""
                cumulative_display = ""  # Cumulative yielded text (with <think>)
                in_thinking = False
                has_content_tokens = False
                tool_calls_acc = {}  # Structured delta.tool_calls fragments
                has_structured_tc = False
                _iter_usage = None
                _iter_timings = None
                _iter_finish_reason = None
                _stream_done = False
                _last_emitted = ""
                provisional_render_html_tool_call_ids = set()
                _suppress_visible_output = _forced_tool_call_pending

                stream_timeout = httpx.Timeout(
                    connect = 10,
                    read = 0.5,
                    write = 10,
                    pool = 10,
                )
                with httpx.Client(
                    timeout = stream_timeout,
                    limits = httpx.Limits(max_keepalive_connections = 0),
                ) as client:
                    with self._stream_with_retry(
                        client,
                        url,
                        payload,
                        cancel_event,
                        headers = _auth_headers,
                    ) as response:
                        if response.status_code != 200:
                            error_body = response.read().decode()
                            raise RuntimeError(
                                f"llama-server returned {response.status_code}: " f"{error_body}"
                            )

                        raw_buf = ""
                        for raw_chunk in self._iter_text_cancellable(
                            response,
                            cancel_event,
                        ):
                            raw_buf += raw_chunk
                            while "\n" in raw_buf:
                                line, raw_buf = raw_buf.split("\n", 1)
                                line = line.strip()

                                if not line:
                                    continue
                                if line == "data: [DONE]":
                                    # Flush thinking state for STREAMING
                                    if detect_state == _S_STREAMING and in_thinking:
                                        if has_content_tokens:
                                            cumulative_display += "</think>"
                                            if not _suppress_visible_output:
                                                yield {
                                                    "type": "content",
                                                    "text": _strip_tool_markup(
                                                        cumulative_display,
                                                        final = True,
                                                    ),
                                                }
                                        else:
                                            cumulative_display = reasoning_accum
                                            if not _suppress_visible_output:
                                                yield {
                                                    "type": "content",
                                                    "text": cumulative_display,
                                                }
                                    _stream_done = True
                                    break  # exit inner while
                                if not line.startswith("data: "):
                                    continue

                                try:
                                    chunk_data = json.loads(line[6:])
                                    _ct = chunk_data.get("timings")
                                    if _ct:
                                        _iter_timings = _ct
                                    _cu = chunk_data.get("usage")
                                    if _cu:
                                        _iter_usage = _cu

                                    choices = chunk_data.get("choices", [])
                                    if not choices:
                                        continue

                                    delta = choices[0].get("delta", {})
                                    _fr = choices[0].get("finish_reason")
                                    if _fr:
                                        _iter_finish_reason = _fr

                                    # ── Structured tool_calls ──
                                    tc_deltas = delta.get("tool_calls")
                                    if tc_deltas:
                                        # llama-server can emit visible assistant
                                        # preface content before native structured
                                        # tool_calls. Preserve content_accum as
                                        # the assistant pre-tool text and still
                                        # drain/execute the structured call.
                                        has_structured_tc = True
                                        detect_state = _S_DRAINING
                                        for tc_d in tc_deltas:
                                            idx = tc_d.get("index", 0)
                                            if idx not in tool_calls_acc:
                                                tool_calls_acc[idx] = {
                                                    "id": tc_d.get("id", f"call_{idx}"),
                                                    "type": "function",
                                                    "function": {
                                                        "name": "",
                                                        "arguments": "",
                                                    },
                                                }
                                            elif tc_d.get("id"):
                                                # Update ID if a real one
                                                # arrives on a later delta.
                                                tool_calls_acc[idx]["id"] = tc_d["id"]
                                            func = tc_d.get("function", {})
                                            if func.get("name"):
                                                tool_calls_acc[idx]["function"]["name"] += func[
                                                    "name"
                                                ]
                                            if func.get("arguments"):
                                                tool_calls_acc[idx]["function"]["arguments"] += (
                                                    func["arguments"]
                                                )
                                            current_name = tool_calls_acc[idx]["function"].get(
                                                "name", ""
                                            )
                                            fallback_id = f"call_{idx}"
                                            current_id = tool_calls_acc[idx].get("id", fallback_id)
                                            already_started = (
                                                current_id in provisional_render_html_tool_call_ids
                                            )
                                            has_real_id = current_id != fallback_id
                                            if (
                                                current_name == "render_html"
                                                and not _tool_succeeded("render_html")
                                                and any(
                                                    (
                                                        (tool.get("function") or {}).get("name")
                                                        == "render_html"
                                                    )
                                                    for tool in active_tools
                                                )
                                                and not already_started
                                                and not provisional_render_html_tool_call_ids
                                                and has_real_id
                                            ):
                                                provisional_render_html_tool_call_ids.add(
                                                    current_id
                                                )
                                                yield {
                                                    "type": "tool_start",
                                                    "tool_name": "render_html",
                                                    "tool_call_id": current_id,
                                                    "arguments": {},
                                                    "provenance": tool_event_provenance(
                                                        provisional = True,
                                                    ),
                                                }
                                        continue

                                    # ── Reasoning tokens ──
                                    # Yield only in STREAMING. In BUFFERING and
                                    # DRAINING, accumulate silently so we don't
                                    # corrupt the consumer's prev_text tracker
                                    # (routes/inference.py never resets it
                                    # between tool iterations).
                                    reasoning = delta.get("reasoning_content", "")
                                    if reasoning:
                                        reasoning_accum += reasoning
                                        if detect_state == _S_STREAMING:
                                            if not in_thinking:
                                                cumulative_display += "<think>"
                                                in_thinking = True
                                            cumulative_display += reasoning
                                            if not _suppress_visible_output:
                                                yield {
                                                    "type": "content",
                                                    "text": cumulative_display,
                                                }

                                    # ── Content tokens ──
                                    token = delta.get("content", "")
                                    if token:
                                        has_content_tokens = True
                                        content_accum += token

                                        if detect_state == _S_DRAINING:
                                            pass  # accumulate silently

                                        elif detect_state == _S_STREAMING:
                                            if in_thinking:
                                                cumulative_display += "</think>"
                                                in_thinking = False
                                            cumulative_display += token
                                            cleaned = _strip_tool_markup_streaming(
                                                cumulative_display
                                            )
                                            if len(cleaned) > len(_last_emitted):
                                                _last_emitted = cleaned
                                                if not _suppress_visible_output:
                                                    yield {
                                                        "type": "content",
                                                        "text": cleaned,
                                                    }

                                        elif detect_state == _S_BUFFERING:
                                            content_buffer += token
                                            stripped_buf = content_buffer.lstrip()
                                            if not stripped_buf:
                                                continue

                                            # Check tool signal prefixes.
                                            is_prefix = False
                                            is_match = False
                                            for sig in _tool_xml_signals:
                                                if stripped_buf.startswith(sig):
                                                    is_match = True
                                                    break
                                                if sig.startswith(stripped_buf):
                                                    is_prefix = True
                                                    break

                                            if is_match:
                                                # Tool signal -- flush any visible
                                                # prefix before DRAINING so the
                                                # route sends it before tool_start.
                                                if reasoning_accum:
                                                    cumulative_display += "<think>"
                                                    cumulative_display += reasoning_accum
                                                    cumulative_display += "</think>"
                                                cumulative_display += content_buffer
                                                cleaned = _strip_tool_markup_streaming(
                                                    cumulative_display,
                                                    force = True,
                                                )
                                                if len(cleaned) > len(_last_emitted):
                                                    _last_emitted = cleaned
                                                    if not _suppress_visible_output:
                                                        yield {
                                                            "type": "content",
                                                            "text": cleaned,
                                                        }
                                                detect_state = _S_DRAINING
                                            elif (
                                                is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS
                                            ):
                                                pass  # keep buffering
                                            else:
                                                # Not a tool -- flush buffer
                                                detect_state = _S_STREAMING
                                                # Flush reasoning accumulated
                                                # during BUFFERING.
                                                if reasoning_accum:
                                                    cumulative_display += "<think>"
                                                    cumulative_display += reasoning_accum
                                                    cumulative_display += "</think>"
                                                cumulative_display += content_buffer
                                                cleaned = _strip_tool_markup(
                                                    cumulative_display,
                                                )
                                                if len(cleaned) > len(_last_emitted):
                                                    _last_emitted = cleaned
                                                    if not _suppress_visible_output:
                                                        yield {
                                                            "type": "content",
                                                            "text": cleaned,
                                                        }

                                except json.JSONDecodeError:
                                    logger.debug(f"Skipping malformed SSE line: {line[:100]}")
                            if _stream_done:
                                break  # exit outer for

                # ── Resolve BUFFERING at stream end ──
                if detect_state == _S_BUFFERING:
                    stripped_buf = content_buffer.lstrip()
                    if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals):
                        detect_state = _S_DRAINING
                    elif content_accum or reasoning_accum:
                        detect_state = _S_STREAMING
                        if content_buffer:
                            # Flush reasoning first.
                            if reasoning_accum:
                                cumulative_display += "<think>"
                                cumulative_display += reasoning_accum
                                cumulative_display += "</think>"
                            cumulative_display += content_buffer
                            if not _suppress_visible_output:
                                yield {
                                    "type": "content",
                                    "text": _strip_tool_markup(
                                        cumulative_display,
                                        final = True,
                                    ),
                                }
                        elif reasoning_accum and not has_content_tokens:
                            # Reasoning-only response: show reasoning as plain
                            # text, matching the final streaming pass for
                            # models that put everything in reasoning.
                            cumulative_display = reasoning_accum
                            if not _suppress_visible_output:
                                yield {
                                    "type": "content",
                                    "text": cumulative_display,
                                }
                    else:
                        return

                # ── STREAMING path: no tool call ──
                if detect_state == _S_STREAMING:
                    # Safety net: check for XML tool signals in content. The
                    # route layer resets prev_text on tool_start, so post-tool
                    # synthesis streams correctly even if content was emitted
                    # before the tool XML.
                    _safety_tc = None
                    if any(s in content_accum for s in _tool_xml_signals):
                        _safety_tc = self._parse_tool_calls_from_text(
                            content_accum,
                            allow_incomplete = auto_heal_tool_calls,
                        )
                    if not _safety_tc:
                        # ── Re-prompt on plan-without-action ──
                        # If the model described its intent (forward-looking
                        # language) without calling a tool, nudge it to act.
                        # Fires at most once per request, only on short
                        # responses with intent signals -- "4" or "Hello!"
                        # won't trigger it. Use content if available, else
                        # fall back to reasoning text (reasoning-only stalls).
                        _stripped = content_accum.strip()
                        if not _stripped:
                            _stripped = reasoning_accum.strip()
                        _render_html_already_done_intent = _tool_succeeded(
                            "render_html"
                        ) and re.search(
                            r"(?i)\brender[_\s-]?html\b",
                            _stripped,
                        )
                        if (
                            auto_heal_tool_calls
                            and active_tools
                            and not _render_html_already_done_intent
                            and _reprompt_count < _MAX_REPROMPTS
                            and _is_short_intent_without_action(_stripped)
                        ):
                            _reprompt_count += 1
                            logger.info(
                                f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: "
                                f"model responded without calling tools "
                                f"({len(_stripped)} chars)"
                            )
                            conversation.append(
                                {
                                    "role": "assistant",
                                    "content": _stripped,
                                }
                            )
                            available_tool_names = [
                                (tool.get("function") or {}).get("name")
                                for tool in active_tools
                                if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
                            ]
                            available_tool_names = [name for name in available_tool_names if name]
                            tool_hint = " or ".join(available_tool_names) or "an available tool"
                            _forced_tool_call_pending = True
                            conversation.append(
                                {
                                    "role": "user",
                                    "content": (
                                        "You have access to enabled tools. If a tool is needed to satisfy "
                                        "the user's request or complete the action you described, call "
                                        f"{tool_hint} now. If no tool is needed, provide the final answer "
                                        "and follow the user's requested format."
                                    ),
                                }
                            )
                            # Accumulate tokens and timing from this iteration.
                            _fu_r = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
                            _accumulated_completion_tokens += _fu_r.get("completion_tokens", 0)
                            _it_r = _iter_timings or {}
                            _accumulated_predicted_ms += _it_r.get("predicted_ms", 0)
                            _accumulated_predicted_n += _it_r.get("predicted_n", 0)
                            yield {"type": "status", "text": ""}
                            continue

                        if _forced_tool_call_pending:
                            _forced_tool_call_pending = False
                            if not _should_suppress_forced_no_tool_output(_stripped):
                                if cumulative_display:
                                    forced_visible_text = _strip_tool_markup(
                                        cumulative_display,
                                        final = True,
                                    )
                                elif content_accum:
                                    forced_visible_text = _strip_tool_markup(
                                        content_accum,
                                        final = True,
                                    )
                                else:
                                    forced_visible_text = reasoning_accum
                                if forced_visible_text:
                                    yield {
                                        "type": "content",
                                        "text": forced_visible_text,
                                    }

                        # Content was already streamed.  Yield metadata.
                        yield {"type": "status", "text": ""}
                        _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
                        _fc = _fu.get("completion_tokens", 0)
                        _fp = _fu.get("prompt_tokens", 0)
                        _tc = _fc + _accumulated_completion_tokens
                        if _iter_usage or _iter_timings or _accumulated_completion_tokens:
                            _mt = dict(_iter_timings) if _iter_timings else {}
                            if _accumulated_predicted_ms or _accumulated_predicted_n:
                                _mt["predicted_ms"] = (
                                    _mt.get("predicted_ms", 0) + _accumulated_predicted_ms
                                )
                                _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n
                                _mt["predicted_n"] = _tn
                                _tms = _mt["predicted_ms"]
                                if _tms > 0:
                                    _mt["predicted_per_second"] = _tn / (_tms / 1000.0)
                            yield {
                                "type": "metadata",
                                "usage": {
                                    "prompt_tokens": _fp,
                                    "completion_tokens": _tc,
                                    "total_tokens": _fp + _tc,
                                },
                                "timings": _mt,
                                "finish_reason": _iter_finish_reason,
                            }
                        return

                    # Safety net caught tool XML -- treat as tool call.
                    tool_calls = _safety_tc
                    content_text = _strip_tool_markup(
                        content_accum,
                        final = True,
                        force = True,
                    )
                    logger.info(
                        f"Safety net: parsed {len(tool_calls)} tool call(s) "
                        f"from streamed content"
                    )
                else:
                    # ── DRAINING path: assemble tool_calls ──
                    tool_calls = None
                    content_text = content_accum
                    if has_structured_tc:
                        # Drop incomplete fragments (e.g. from max_tokens
                        # truncation or disconnect).
                        tool_calls = [
                            tool_calls_acc[i]
                            for i in sorted(tool_calls_acc)
                            if (tool_calls_acc[i].get("function", {}).get("name", "").strip())
                        ] or None
                    if not tool_calls and any(s in content_accum for s in _tool_xml_signals):
                        tool_calls = self._parse_tool_calls_from_text(
                            content_accum,
                            allow_incomplete = auto_heal_tool_calls,
                        )
                    if tool_calls and not has_structured_tc:
                        content_text = _strip_tool_markup(
                            content_text,
                            final = True,
                            force = True,
                        )
                    if tool_calls:
                        logger.info(
                            f"Parsed {len(tool_calls)} tool call(s) from "
                            f"{'structured delta' if has_structured_tc else 'content text'}"
                        )
                    if not tool_calls:
                        # DRAINING but no tool calls (false positive). Merge
                        # accumulated metrics from prior tool iterations so
                        # they aren't silently dropped.
                        yield {"type": "status", "text": ""}
                        if content_accum:
                            # Strip leaked tool-call XML before yielding.
                            content_accum = _strip_tool_markup(content_accum, final = True)
                        if content_accum:
                            yield {"type": "content", "text": content_accum}
                        _fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
                        _fc = _fu.get("completion_tokens", 0)
                        _fp = _fu.get("prompt_tokens", 0)
                        _tc = _fc + _accumulated_completion_tokens
                        if _iter_usage or _iter_timings or _accumulated_completion_tokens:
                            _mt = dict(_iter_timings) if _iter_timings else {}
                            if _accumulated_predicted_ms or _accumulated_predicted_n:
                                _mt["predicted_ms"] = (
                                    _mt.get("predicted_ms", 0) + _accumulated_predicted_ms
                                )
                                _tn = _mt.get("predicted_n", 0) + _accumulated_predicted_n
                                _mt["predicted_n"] = _tn
                                _tms = _mt["predicted_ms"]
                                if _tms > 0:
                                    _mt["predicted_per_second"] = _tn / (_tms / 1000.0)
                            yield {
                                "type": "metadata",
                                "usage": {
                                    "prompt_tokens": _fp,
                                    "completion_tokens": _tc,
                                    "total_tokens": _fp + _tc,
                                },
                                "timings": _mt,
                                "finish_reason": _iter_finish_reason,
                            }
                        return

                # ── Execute tool calls ──
                _accumulated_completion_tokens += (
                    _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
                ).get("completion_tokens", 0)
                _it = _iter_timings or {}
                _accumulated_predicted_ms += _it.get("predicted_ms", 0)
                _accumulated_predicted_n += _it.get("predicted_n", 0)

                # disable_parallel_tool_use: execute only the first tool call
                # this turn. Truncate before building assistant_msg so the
                # conversation stays consistent and extra calls are never executed.
                if disable_parallel_tool_use and tool_calls and len(tool_calls) > 1:
                    tool_calls = tool_calls[:1]

                assistant_msg: dict = {"role": "assistant", "content": content_text}
                assistant_appended = False

                for tc in tool_calls or []:
                    func = tc.get("function", {})
                    tool_name = func.get("name", "")
                    provisional_render_html_match = (
                        tool_name == "render_html"
                        and tc.get("id") in provisional_render_html_tool_call_ids
                    )
                    decision = tool_controller.prepare_call(
                        tc,
                        forced = _forced_tool_call_pending,
                        provisional = provisional_render_html_match,
                    )

                    if not decision.should_execute:
                        if content_text and not assistant_appended:
                            conversation.append(assistant_msg)
                            assistant_appended = True
                        completion = tool_controller.record_noop(decision)
                        conversation.append(completion.model_message())
                        if _forced_tool_call_pending:
                            _forced_tool_call_pending = False
                        logger.info(
                            "Suppressed local GGUF tool call as internal no-op: "
                            f"action={decision.action} tool={decision.tool_name}"
                        )
                        break

                    if not assistant_appended:
                        assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
                        conversation.append(assistant_msg)
                        assistant_appended = True
                    else:
                        assistant_msg.setdefault("tool_calls", []).append(
                            decision.as_assistant_tool_call()
                        )

                    yield {"type": "status", "text": decision.status_text}
                    yield decision.tool_start_event()

                    _effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
                    # RAG: cap paraphrased KB re-searches that slip past the dup guard.
                    if (
                        decision.tool_name == "search_knowledge_base"
                        and _kb_search_count >= RAG_MAX_SEARCHES_PER_TURN
                    ):
                        result = RAG_SEARCH_CAP_NUDGE
                    else:
                        result = execute_tool(
                            decision.tool_name,
                            decision.arguments,
                            cancel_event = cancel_event,
                            timeout = _effective_timeout,
                            session_id = session_id,
                            rag_scope = rag_scope,
                        )
                        if decision.tool_name == "search_knowledge_base":
                            _kb_search_count += 1
                    completion = tool_controller.record_result(decision, result)
                    yield completion.tool_end_event()
                    conversation.append(completion.tool_message())

                    if _forced_tool_call_pending:
                        _forced_tool_call_pending = False

                # Clear tool status badge before next generation/final pass.
                yield {"type": "status", "text": ""}
                if tool_controller.force_final_answer or not tool_controller.active_tools():
                    _append_budget_exhausted_nudge = False
                    break
                continue

            except httpx.ConnectError:
                raise RuntimeError("Lost connection to llama-server")
            except Exception as e:
                if cancel_event is not None and cancel_event.is_set():
                    return
                raise

        # ── Tool iteration cap reached -- synthesize final answer ──
        # The model used all iterations without a final text response. Nudge
        # the final streaming pass to produce a useful answer instead of
        # continuing to request tools.
        if max_tool_iterations > 0 and _append_budget_exhausted_nudge:
            conversation.append(
                {
                    "role": "user",
                    "content": (
                        "You have used all available tool calls. Based on "
                        "everything you have found so far, provide your final "
                        "answer now. Do not call any more tools."
                    ),
                }
            )

        # Clear status.
        yield {"type": "status", "text": ""}

        # Final streaming pass with the full conversation context.
        stream_payload = {
            "messages": conversation,
            "stream": True,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k if top_k >= 0 else 0,
            "min_p": min_p,
            "repeat_penalty": repetition_penalty,
            "presence_penalty": presence_penalty,
        }
        _reasoning_kw = self._request_reasoning_kwargs(
            enable_thinking, reasoning_effort, preserve_thinking
        )
        if _reasoning_kw is not None:
            stream_payload["chat_template_kwargs"] = _reasoning_kw
        stream_payload["max_tokens"] = (
            max_tokens
            if max_tokens is not None
            else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
        )
        stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
        if stop:
            stream_payload["stop"] = stop
        if seed is not None:
            stream_payload["seed"] = seed
        stream_payload["stream_options"] = {"include_usage": True}

        cumulative = ""
        _last_emitted = ""
        in_thinking = False
        has_content_tokens = False
        reasoning_text = ""
        _metadata_usage = None
        _metadata_timings = None
        _metadata_finish_reason = None
        _stream_done = False

        try:
            stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
            _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
            with httpx.Client(
                timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
            ) as client:
                with self._stream_with_retry(
                    client,
                    url,
                    stream_payload,
                    cancel_event,
                    headers = _auth_headers,
                ) as response:
                    if response.status_code != 200:
                        error_body = response.read().decode()
                        raise RuntimeError(
                            f"llama-server returned {response.status_code}: {error_body}"
                        )

                    buffer = ""
                    for raw_chunk in self._iter_text_cancellable(response, cancel_event):
                        buffer += raw_chunk
                        while "\n" in buffer:
                            line, buffer = buffer.split("\n", 1)
                            line = line.strip()

                            if not line:
                                continue
                            if line == "data: [DONE]":
                                if in_thinking:
                                    if has_content_tokens:
                                        cumulative += "</think>"
                                        yield {
                                            "type": "content",
                                            "text": _strip_tool_markup(cumulative, final = True),
                                        }
                                    else:
                                        cumulative = reasoning_text
                                        yield {"type": "content", "text": cumulative}
                                _stream_done = True
                                break  # exit inner while
                            if not line.startswith("data: "):
                                continue

                            try:
                                chunk_data = json.loads(line[6:])
                                # Capture server timings/usage from final chunks.
                                _chunk_timings = chunk_data.get("timings")
                                if _chunk_timings:
                                    _metadata_timings = _chunk_timings
                                _chunk_usage = chunk_data.get("usage")
                                if _chunk_usage:
                                    _metadata_usage = _chunk_usage
                                choices = chunk_data.get("choices", [])
                                if choices:
                                    delta = choices[0].get("delta", {})
                                    _fr = choices[0].get("finish_reason")
                                    if _fr:
                                        _metadata_finish_reason = _fr

                                    reasoning = delta.get("reasoning_content", "")
                                    if reasoning:
                                        reasoning_text += reasoning
                                        if not in_thinking:
                                            cumulative += "<think>"
                                            in_thinking = True
                                        cumulative += reasoning
                                        yield {"type": "content", "text": cumulative}

                                    token = delta.get("content", "")
                                    if token:
                                        has_content_tokens = True
                                        if in_thinking:
                                            cumulative += "</think>"
                                            in_thinking = False
                                        cumulative += token
                                        cleaned = _strip_tool_markup(cumulative)
                                        # Emit only when cleaned text grows (monotonic).
                                        if len(cleaned) > len(_last_emitted):
                                            _last_emitted = cleaned
                                            yield {"type": "content", "text": cleaned}
                            except json.JSONDecodeError:
                                logger.debug(f"Skipping malformed SSE line: {line[:100]}")
                        if _stream_done:
                            break  # exit outer for
                    _final_usage = _metadata_usage or {}
                    _final_completion = _final_usage.get("completion_tokens", 0)
                    _final_prompt = _final_usage.get("prompt_tokens", 0)
                    _total_completion = _final_completion + _accumulated_completion_tokens
                    if _metadata_usage or _metadata_timings or _metadata_finish_reason:
                        _merged_timings = dict(_metadata_timings) if _metadata_timings else {}
                        if _accumulated_predicted_ms or _accumulated_predicted_n:
                            _merged_timings["predicted_ms"] = (
                                _merged_timings.get("predicted_ms", 0) + _accumulated_predicted_ms
                            )
                            _total_predicted_n = (
                                _merged_timings.get("predicted_n", 0) + _accumulated_predicted_n
                            )
                            _merged_timings["predicted_n"] = _total_predicted_n
                            _total_predicted_ms = _merged_timings["predicted_ms"]
                            if _total_predicted_ms > 0:
                                _merged_timings["predicted_per_second"] = _total_predicted_n / (
                                    _total_predicted_ms / 1000.0
                                )
                        yield {
                            "type": "metadata",
                            "usage": {
                                "prompt_tokens": _final_prompt,
                                "completion_tokens": _total_completion,
                                "total_tokens": _final_prompt + _total_completion,
                            },
                            "timings": _merged_timings,
                            "finish_reason": _metadata_finish_reason,
                        }

        except httpx.ConnectError:
            raise RuntimeError("Lost connection to llama-server")
        except Exception as e:
            if cancel_event is not None and cancel_event.is_set():
                return
            raise

    # ── Prompt token counting ──────────────────────────────────

    def count_chat_tokens(
        self,
        messages,
        system = None,
        tools = None,
        strict: bool = False,
    ) -> int:
        """Count prompt tokens for a chat request via llama-server.

        Non-strict callers keep the historical best-effort behavior and receive
        0 when a count cannot be determined. Strict callers (public count_tokens
        endpoints) get an exception instead of a successful-looking zero when
        tokenizer/template calls fail or a multimodal prompt would fall back to a
        text-only approximation.
        """
        if not self.is_loaded:
            if strict:
                raise RuntimeError("llama-server is not loaded")
            return 0

        def _has_non_text_content(content) -> bool:
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, str):
                        continue
                    if not isinstance(block, dict):
                        return True
                    if block.get("type") == "text" and isinstance(block.get("text"), str):
                        continue
                    if isinstance(block.get("text"), str):
                        continue
                    return True
            return False

        def _has_non_text_prompt_parts() -> bool:
            if _has_non_text_content(system):
                return True
            for msg in messages or []:
                if isinstance(msg, dict) and _has_non_text_content(msg.get("content", "")):
                    return True
            return False

        def _block_text(content) -> str:
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                parts = []
                for block in content:
                    if isinstance(block, dict):
                        if block.get("type") == "text" and isinstance(block.get("text"), str):
                            parts.append(block["text"])
                        elif isinstance(block.get("text"), str):
                            parts.append(block["text"])
                    elif isinstance(block, str):
                        parts.append(block)
                return "".join(parts)
            return ""

        # Normalize system into a leading message / plain text.
        system_text = ""
        if isinstance(system, str):
            system_text = system
        elif isinstance(system, list):
            system_text = _block_text(system)

        try:
            _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
            with httpx.Client(timeout = 10, headers = _auth_headers) as client:

                def _tokenize(text: str) -> int:
                    r = client.post(
                        f"{self.base_url}/tokenize",
                        json = {"content": text, "add_special": True},
                    )
                    if r.status_code != 200:
                        if strict:
                            raise RuntimeError("llama-server tokenizer failed")
                        return 0
                    tokens = r.json().get("tokens", [])
                    if not isinstance(tokens, list):
                        if strict:
                            raise RuntimeError("llama-server tokenizer returned invalid tokens")
                        return 0
                    return len(tokens)

                # 1. Try /apply-template to render the real chat prompt.
                template_messages = list(messages) if messages else []
                if system_text:
                    template_messages = [
                        {"role": "system", "content": system_text}
                    ] + template_messages
                apply_template_failed = False
                try:
                    # llama-server's /apply-template renders tool declarations
                    # into the prompt when ``tools`` is supplied, so pass them
                    # through — otherwise tool-schema tokens go uncounted.
                    template_body = {"messages": template_messages}
                    if tools:
                        template_body["tools"] = tools
                    resp = client.post(
                        f"{self.base_url}/apply-template",
                        json = template_body,
                    )
                    if resp.status_code == 200:
                        prompt = resp.json().get("prompt", "")
                        if isinstance(prompt, str):
                            return _tokenize(prompt)
                    apply_template_failed = True
                except Exception:
                    apply_template_failed = True

                if strict and apply_template_failed and _has_non_text_prompt_parts():
                    raise RuntimeError(
                        "cannot fall back to text-only token counting for multimodal messages"
                    )

                # 2. Fallback: concatenate plain text and tokenize. Append a
                # serialized form of the tools so they still contribute to the
                # count when /apply-template is unavailable.
                parts = []
                if system_text:
                    parts.append(system_text)
                for msg in messages or []:
                    if isinstance(msg, dict):
                        parts.append(_block_text(msg.get("content", "")))
                if tools:
                    try:
                        parts.append(json.dumps(tools, ensure_ascii = False))
                    except Exception:
                        pass
                return _tokenize("\n".join(p for p in parts if p))
        except Exception:
            if strict:
                raise
            return 0

    # ── TTS support ────────────────────────────────────────────

    def detect_audio_type(self) -> Optional[str]:
        """Detect audio/TTS codec; swallows errors (use _strict to distinguish)."""
        try:
            return self._detect_audio_type_strict()
        except Exception as e:
            logger.debug(f"Audio type detection failed: {e}")
            return None

    def _detect_audio_type_strict(self) -> Optional[str]:
        """Codec name on match, None on non-audio, raises on transport/JSON errors."""
        if not self.is_loaded:
            return None
        _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
        with httpx.Client(timeout = 10, headers = _auth_headers) as client:

            def _detok(tid: int) -> str:
                # Non-200 means "marker not in vocab" -- keep probing.
                # Transport / JSON errors still raise.
                r = client.post(f"{self.base_url}/detokenize", json = {"tokens": [tid]})
                if r.status_code != 200:
                    return ""
                return r.json().get("content", "")

            def _tok(text: str) -> list[int]:
                r = client.post(
                    f"{self.base_url}/tokenize",
                    json = {"content": text, "add_special": False},
                )
                if r.status_code != 200:
                    return []
                return r.json().get("tokens", [])

            # Codec-specific tokens (not generic ones that non-audio models may have)
            if "<custom_token_" in _detok(128258) and "<custom_token_" in _detok(128259):
                return "snac"
            if len(_tok("<|AUDIO|>")) == 1 and len(_tok("<|audio_eos|>")) == 1:
                return "csm"
            if len(_tok("<|startoftranscript|>")) == 1:
                return "whisper"
            # Gemma 3n: <audio_soft_token>; Gemma 4: <|audio|> (not csm's <|AUDIO|>).
            if len(_tok("<audio_soft_token>")) == 1 or len(_tok("<|audio|>")) == 1:
                return "audio_vlm"
            if len(_tok("<|bicodec_semantic_0|>")) == 1 and len(_tok("<|bicodec_global_0|>")) == 1:
                return "bicodec"
            if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1:
                return "dac"
        return None

    # Prompt format per codec: (template, stop_tokens, needs_token_ids).
    # Matches InferenceBackend._generate_snac/bicodec/dac.
    _TTS_PROMPTS = {
        "snac": (
            "<custom_token_3>{text}<|eot_id|><custom_token_4>",
            ["<custom_token_2>"],
            True,
        ),
        "bicodec": (
            "<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>",
            ["<|im_end|>", "</s>"],
            False,
        ),
        "dac": (
            "<|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|><|global_features_start|>\n",
            ["<|im_end|>", "<|audio_end|>"],
            False,
        ),
    }

    _codec_mgr = None  # Shared AudioCodecManager instance

    def init_audio_codec(self, audio_type: str) -> None:
        """Load the audio codec at model load time (mirrors the non-GGUF path)."""
        import torch
        from core.inference.audio_codecs import AudioCodecManager

        if LlamaCppBackend._codec_mgr is None:
            LlamaCppBackend._codec_mgr = AudioCodecManager()

        device = "cuda" if torch.cuda.is_available() else "cpu"
        model_repo_path = None

        # BiCodec needs a repo with BiCodec/ weights -- download canonical SparkTTS
        if audio_type == "bicodec":
            from huggingface_hub import snapshot_download
            import os

            repo_path = snapshot_download("unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B")
            model_repo_path = os.path.abspath(repo_path)

        LlamaCppBackend._codec_mgr.load_codec(audio_type, device, model_repo_path = model_repo_path)
        logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}")

    def generate_audio_response(
        self,
        text: str,
        audio_type: str,
        temperature: float = 0.6,
        top_p: float = 0.95,
        top_k: int = 50,
        min_p: float = 0.0,
        max_new_tokens: int = 2048,
        repetition_penalty: float = 1.1,
    ) -> tuple:
        """
        Generate TTS audio via llama-server /completion + codec decode.
        Returns (wav_bytes, sample_rate).
        """
        if audio_type not in self._TTS_PROMPTS:
            raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")

        tpl, stop, need_ids = self._TTS_PROMPTS[audio_type]

        payload: dict = {
            "prompt": tpl.format(text = text),
            "stream": False,
            "n_predict": max_new_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k if top_k >= 0 else 0,
            "min_p": min_p,
            "repeat_penalty": repetition_penalty,
        }
        if stop:
            payload["stop"] = stop
        if need_ids:
            payload["n_probs"] = 1

        _auth_headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
        with httpx.Client(timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers) as client:
            resp = client.post(f"{self.base_url}/completion", json = payload)
            if resp.status_code != 200:
                raise RuntimeError(f"llama-server returned {resp.status_code}: {resp.text}")

        data = resp.json()
        token_ids = (
            [p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
            if need_ids
            else None
        )

        import torch

        device = "cuda" if torch.cuda.is_available() else "cpu"
        return LlamaCppBackend._codec_mgr.decode(
            audio_type, device, token_ids = token_ids, text = data.get("content", "")
        )