diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 9dd56489eb..b13cd1c851 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -143,6 +143,17 @@ async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depend ) +async def authenticated_via_api_key( + credentials: HTTPAuthorizationCredentials = Depends(security), +) -> bool: + """True when the caller used an sk-unsloth API key, not a UI session JWT. + + Lets routes treat programmatic API callers differently from the Studio UI + (e.g. refuse a teardown the UI would allow). + """ + return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX)) + + async def get_current_subject_allow_password_change( credentials: HTTPAuthorizationCredentials = Depends(security), ) -> str: diff --git a/studio/backend/core/inference/llama_keepwarm.py b/studio/backend/core/inference/llama_keepwarm.py new file mode 100644 index 0000000000..4ce663c3ce --- /dev/null +++ b/studio/backend/core/inference/llama_keepwarm.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in idle auto-unload (TTL keep-warm) for the local llama.cpp model. + +Off by default (idle seconds = 0). When enabled, a background loop unloads the +loaded GGUF once it has been idle for the configured TTL, freeing VRAM. A +pure-ASGI middleware tracks in-flight inference requests so a long stream that +outlives the TTL is never unloaded mid-response. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading +import time + +from loggers import get_logger + +logger = get_logger(__name__) + +_lock = threading.Lock() +_inflight = 0 +# Requests blocked on the unload gate but not yet counted in _inflight: the idle +# loop must not unload while one is waiting (it would unload out from under it). +_pending = 0 +_last_active = time.monotonic() +# The (id, quant) idle-unload last freed, so an alias/unknown request that would +# otherwise 503 against an empty backend can reload it (set on unload, cleared on +# reload). Storing the quant means the reload restores the exact freed variant. +_last_unloaded_model = None +# Guards inflight bumps against the idle-check-then-unload race, and blocks new +# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is +# shared across every event loop in the process, so a per-loop gate would let a +# request on loop B start inference while a swap on loop A tears the model down. +_lifecycle_lock = threading.Lock() + + +@contextlib.asynccontextmanager +async def _unload_gate(): + # Acquire off the loop: non-blocking first (the common uncontended case), else + # poll a non-blocking acquire off a short sleep. Polling keeps the wait off this + # loop AND cancellation-safe -- a cancel lands during the sleep, when the gate is + # not held, so it never leaks (mirrors the auto-switch swap gate). + while not _lifecycle_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + try: + yield + finally: + _lifecycle_lock.release() + + +_INFERENCE_PREFIXES = ("/v1/", "/api/inference/") +_INFERENCE_SUFFIXES = ( + "/chat/completions", + "/completions", + "/messages", + "/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages + "/embeddings", + "/responses", + "/generate/stream", # Studio's own streaming route on the same llama-server + "/audio/generate", # direct GGUF TTS; can outlive the idle TTL +) + + +def _is_inference_path(path: str) -> bool: + if path.startswith(_INFERENCE_PREFIXES) and path.endswith(_INFERENCE_SUFFIXES): + return True + # Public checkpoint preview (/p/{run}/v1/chat/completions) delegates to the + # chat handler and streams from the same backend, so protect it from idle unload. + return path.startswith("/p/") and path.endswith("/v1/chat/completions") + + +def _note_pending() -> None: + global _pending + with _lock: + _pending += 1 + + +def _note_unpending() -> None: + global _pending + with _lock: + _pending = max(0, _pending - 1) + + +def _note_start() -> None: + # Do not stamp _last_active here: while _inflight > 0 the model is already + # protected (see _is_idle), and stamping on start lets an external-provider + # request that is later untracked still reset the local idle timer. + global _inflight, _pending + with _lock: + _pending = max(0, _pending - 1) + _inflight += 1 + + +def _note_end() -> None: + global _inflight, _last_active + with _lock: + _inflight = max(0, _inflight - 1) + _last_active = time.monotonic() + + +def _note_untracked_end() -> None: + # Drop a request that never used the local GGUF without stamping local + # activity, so periodic external-provider traffic can't keep the model warm. + global _inflight + with _lock: + _inflight = max(0, _inflight - 1) + + +def _is_idle(ttl_seconds: float) -> bool: + with _lock: + return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds + + +def _note_activity() -> None: + """Stamp activity, e.g. on a (re)load, so the model survives at least one TTL.""" + global _last_active + with _lock: + _last_active = time.monotonic() + + +def other_inference_request_count( + current_request_counted: bool = True, *, include_pending: bool = True +) -> int: + """Tracked inference requests other than the current route call. + + The middleware counts OpenAI-compatible requests before route code runs, so + the caller is excluded by default. Idle-unload counts pending waiters too (a + swap holding the gate would unload out from under them). The swap guard passes + include_pending=False: a pending request is blocked in the middleware and has + not started inference, so it can't be the request a swap would interrupt. + """ + with _lock: + active = _inflight + if current_request_counted and active > 0: + active -= 1 + return max(0, active) + (_pending if include_pending else 0) + + +# Set on the ASGI scope by a route that proved this request won't touch +# llama.cpp (e.g. it proxied to an external provider), so the keep-warm count +# excludes it and the middleware skips its own end-decrement. +_UNTRACKED_SCOPE_KEY = "_unsloth_keepwarm_untracked" + + +def untrack_current_request(scope) -> None: + """Drop this request from the in-flight count once the route knows it won't + use the local GGUF, so unrelated external-provider traffic can't trip the + swap busy guard. Idempotent; the middleware then skips its end-decrement.""" + if not isinstance(scope, dict) or scope.get(_UNTRACKED_SCOPE_KEY): + return + scope[_UNTRACKED_SCOPE_KEY] = True + _note_untracked_end() + + +def inference_lifecycle_gate(): + """The gate a model swap holds so new inference can't start mid-load. Process- + wide, so a swap on one loop blocks inference starting on any other loop.""" + return _unload_gate() + + +def note_model_loaded() -> None: + """Record a successful GGUF load: stamp activity and drop any reload stash so + a manual load clears it synchronously, not only on the next idle poll.""" + _note_activity() + _set_last_unloaded(None) + + +def note_model_unloaded() -> None: + """Record a deliberate (user/API) unload: drop any idle reload stash so the next + request can't resurrect the just-unloaded model. The idle loop unloads via the + backend directly and then stashes the freed model for an alias reload; an + explicit unload instead means "stay unloaded", so it must not stamp activity.""" + _set_last_unloaded(None) + + +def get_last_unloaded_model(): + with _lock: + return _last_unloaded_model + + +def _set_last_unloaded(value) -> None: + global _last_unloaded_model + with _lock: + _last_unloaded_model = value + + +class LlamaKeepWarmMiddleware: + """Pure ASGI: count in-flight inference requests and stamp activity on completion.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + # Inference endpoints are all POST; skipping non-POST avoids counting CORS + # preflight (OPTIONS). ``or ""`` guards an explicit None path. + if ( + scope.get("type") != "http" + or scope.get("method") != "POST" + or not _is_inference_path(scope.get("path") or "") + ): + await self.app(scope, receive, send) + return + # Always track in-flight on inference paths, even when the feature is off, + # so a stream that starts before idle-unload is enabled can't be unloaded + # mid-response if the operator turns it on during that stream. Counting is + # cheap and invisible to clients (the response is proxied unchanged). + # Mark pending before the gate so the idle loop (which holds the gate while + # unloading) can't free the model while this request is waiting to start. + _note_pending() + started = False + try: + async with _unload_gate(): + _note_start() + started = True + finally: + if not started: + _note_unpending() + ended = {"done": False} + status = {"code": None} + + def _finish() -> None: + # A route that untracked itself already decremented; don't double-count. + if ended["done"]: + return + ended["done"] = True + if scope.get(_UNTRACKED_SCOPE_KEY): + return + # This middleware runs before FastAPI auth, so a 401/403 reaches here + # without ever touching llama.cpp. Decrement the in-flight count (to + # balance _note_start) but do NOT stamp activity, or repeated + # unauthenticated probes on an exposed server would keep the model warm + # and never let idle-unload free VRAM. + if status["code"] in (401, 403): + _note_untracked_end() + else: + _note_end() + + async def send_wrapper(message): + if message.get("type") == "http.response.start": + status["code"] = message.get("status") + # Final body frame marks the end of a (possibly streaming) response. + elif message.get("type") == "http.response.body" and not message.get( + "more_body", False + ): + _finish() + await send(message) + + try: + await self.app(scope, receive, send_wrapper) + finally: + _finish() + + +def _loaded_identity(backend): + if not backend.is_loaded or not backend.model_identifier: + return None + # Third slot is the advertised id (repo id) an auto-switch load sets on the + # backend; it's the override key, so an idle stash keyed by the concrete load + # path doesn't drop the user's saved launch flags on the alias reload. + advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier + return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised) + + +async def idle_unload_loop(poll_seconds: float = 15.0) -> None: + """Unload the loaded GGUF once idle past the configured TTL. Inert when off.""" + from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds + + seen_model = None + while True: + await asyncio.sleep(poll_seconds) + try: + ttl = get_auto_unload_idle_seconds() + if ttl <= 0: + continue + from routes.inference import get_llama_cpp_backend + + backend = get_llama_cpp_backend() + # Track by (id, variant): a (re)loaded model -- including the same repo + # at a different quant -- counts as activity so it survives one TTL + # before its first request (loads bypass the activity middleware). + current = _loaded_identity(backend) + if current != seen_model: + seen_model = current + if current is not None: + _note_activity() + _set_last_unloaded(None) # a model is loaded; drop stale stash + async with _unload_gate(): + if backend.is_loaded and _is_idle(ttl): + freed = _loaded_identity(backend) + await asyncio.to_thread(backend.unload_model) + _set_last_unloaded(freed) # let an alias request reload it + logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl) + seen_model = None + except Exception as exc: + logger.debug("idle_unload_loop iteration failed: %s", exc) diff --git a/studio/backend/core/inference/local_model_resolver.py b/studio/backend/core/inference/local_model_resolver.py new file mode 100644 index 0000000000..002cafe2c8 --- /dev/null +++ b/studio/backend/core/inference/local_model_resolver.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve an OpenAI-request ``model`` string to a downloaded local GGUF. + +Used by the opt-in auto-switch path. The match is conservative: only names +that map to an already-downloaded local GGUF (and a quant that is actually on +disk) are eligible, so an arbitrary OpenAI model string still falls through to +the loaded model (drop-in compat) and no surprise multi-GB download is ever +triggered. The local-model scan is cached for a few seconds since auto-switch +consults it per request. +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass +from typing import Optional + +from core.inference.model_ids import public_model_id +from loggers import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen = True) +class _LocalGgufEntry: + loader_id: str # advertised id (repo id / folder name), also the override key + load_path: str # concrete on-disk dir/file passed to /load so it never downloads + variants: tuple[str, ...] # local quant labels; () for a standalone .gguf + + +_CACHE_TTL_S = 5.0 +_lock = threading.Lock() +_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {}) + + +def _is_abs_path_id(value: str) -> bool: + """True when an id is an absolute filesystem path (the ./models and LM Studio + scanners use the on-disk path as the id) rather than a repo id like org/name.""" + from pathlib import Path + try: + return Path(value).is_absolute() + except Exception: + return False + + +def _advertised_loader_id(info) -> Optional[str]: + """The id to advertise for a scanned model: prefer a client-facing alias over + an absolute filesystem path so /v1/models and the override key never expose a + host path (the ./models and LM Studio scanners report the path as info.id).""" + raw_id = getattr(info, "id", None) + if not raw_id or not _is_abs_path_id(raw_id): + return raw_id + for alt in (getattr(info, "model_id", None), getattr(info, "display_name", None)): + if alt and not _is_abs_path_id(alt): + return alt + # No clean alias: strip to a path-free public id so a host path is never advertised. + return public_model_id(raw_id) or raw_id + + +def _resolve_load_dir(p): + """The concrete dir holding the GGUFs. For an HF cache repo (``models--*`` + with ``snapshots/``) this is the latest snapshot dir, so /load takes the + local branch instead of the download-capable repo-id branch.""" + from pathlib import Path + + try: + if (p / "snapshots").is_dir(): + from routes.models import _resolve_hf_cache_realpath + real = _resolve_hf_cache_realpath(p) + if real: + return Path(real) + except Exception: + pass + return p + + +def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]: + """Build an entry only when GGUF quants are on disk (not Transformers/ + safetensors), listing only on-disk quants. ``load_path`` is a concrete local + path so /load resolves the variant locally and never fetches a remote one.""" + from pathlib import Path + from utils.models.model_config import _is_mmproj, list_local_gguf_variants + + path = getattr(info, "path", None) + if not isinstance(path, str): + return None + p = Path(path) + try: + if p.is_file(): + # A standalone .gguf loads by its own path; no quant sub-selection. An + # mmproj companion (vision/audio projector) is not a servable model on + # its own: _scan_models_dir's standalone-file pass does not filter it + # the way the directory scan does, so reject it here or /v1/models would + # advertise a projector and a switch could load it instead of the weights, + # evicting the loaded model. The directory branch below is already mmproj + # free (list_local_gguf_variants drops mmproj quants). + if p.suffix.lower() != ".gguf" or _is_mmproj(p.name): + return None + return _LocalGgufEntry(loader_id, str(p), ()) + load_dir = _resolve_load_dir(p) + variants, _ = list_local_gguf_variants(str(load_dir)) + quants = tuple(v.quant for v in variants if getattr(v, "quant", None)) + return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None + except Exception: + return None + + +def info_has_local_gguf(info) -> bool: + """True when *info* (a LocalModelInfo) points to on-disk GGUF weights the + auto-switch path can load. Read from the files, not ``info.model_format``: the + HF-cache scanner leaves model_format unset for GGUF snapshots, so a + model_format filter would drop every cached GGUF. Lets /v1/models advertise + exactly what /v1 can serve.""" + from pathlib import Path + + path = getattr(info, "path", None) + # Ollama-link entries come from a scanner _build_index intentionally skips (it + # creates symlinks on the request path), so their advertised ids never resolve. + # Don't report them as servable, or /v1/models would list unswitchable models. + if isinstance(path, str) and any( + seg in (".studio_links", "ollama_links") for seg in Path(path).parts + ): + return False + return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None + + +def _build_index() -> dict[str, _LocalGgufEntry]: + """Map normalized id/model_id/display_name -> local GGUF entry. + + Scans the same roots Studio's model picker lists (./models, the active plus + legacy/default HF caches, LM Studio dirs, and user scan folders) so a named + local model is never missed and silently served as the loaded one. Ollama's + scanner is skipped: it creates symlinks as a side effect and this runs on the + request path. + """ + # Lazy import: routes.models imports core.inference, so import at call time. + from pathlib import Path + from routes.models import ( + _scan_models_dir, + _scan_hf_cache, + _scan_lmstudio_dir, + _resolve_hf_cache_dir, + _is_hidden_model, + ) + from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs + + index: dict[str, _LocalGgufEntry] = {} + seen_hf: set[str] = set() + + def _scan_hf_once(directory) -> list: + if directory is None: + return [] + try: + d = Path(directory) + if not d.is_dir(): + return [] + rp = str(d.resolve()) + if rp in seen_hf: + return [] + seen_hf.add(rp) + return _scan_hf_cache(directory) + except Exception as exc: # a missing/malformed root must skip, never crash the index + logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc) + return [] + + # Each source is guarded on its own so one bad root (a permission error, a + # malformed cache) drops only that source, not the whole index. + found: list = [] + try: + found += _scan_models_dir(Path("./models").resolve()) + except Exception as exc: + logger.debug("auto-switch: ./models scan failed: %s", exc) + try: + for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()): + found += _scan_hf_once(hf_dir) + except Exception as exc: + logger.debug("auto-switch: HF cache scan failed: %s", exc) + try: + for lm_dir in lmstudio_model_dirs(): + found += _scan_lmstudio_dir(lm_dir) + except Exception as exc: + logger.debug("auto-switch: LM Studio scan failed: %s", exc) + try: + from storage.studio_db import list_scan_folders + for folder in list_scan_folders(): + try: + fp = Path(folder["path"]) + found += ( + _scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp) + ) + except Exception as exc: + logger.debug("auto-switch: scan folder %r failed: %s", folder, exc) + except Exception as exc: + logger.debug("auto-switch: scan folders enumerate failed: %s", exc) + for info in found: + raw_id = getattr(info, "id", None) + if not raw_id: + continue + # Skip what Studio hides from its pickers (validation probe, RAG embed + # weights): not chat models, so never an auto-switch target. + if _is_hidden_model(raw_id, getattr(info, "path", None)): + continue + # Advertise a client-facing alias, not an absolute filesystem path. + loader_id = _advertised_loader_id(info) + entry = _local_gguf_entry(loader_id, info) + if entry is None: + continue + # Index every alias (including the path) so a client can resolve by any of + # them, even though only the non-path loader_id is advertised. + for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)): + if key: + index.setdefault(key.strip().lower(), entry) + return index + + +def _index() -> dict[str, _LocalGgufEntry]: + global _scan + # Build under the lock so concurrent callers with an expired cache don't all + # run the (multi-dir) scan at once; the rest wait and reuse the fresh result. + with _lock: + now = time.monotonic() + ts, cached = _scan + if now - ts < _CACHE_TTL_S: + return cached + fresh = _build_index() + # Stamp AFTER the scan, not with the pre-scan ``now``: a multi-root scan on + # an install with many local models can itself exceed the TTL, which would + # store the cache already expired and make every request rebuild the index. + _scan = (time.monotonic(), fresh) + return fresh + + +def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]: + """Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None. + + ``load_path`` is the concrete on-disk path to hand /load (so it never fetches + a remote), ``loader_id`` is the advertised id used as the launch-override key. + ``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first + (so ids containing a colon still resolve); else the last ``:VARIANT`` is split + off and resolves only when that quant is on disk. + """ + if not isinstance(requested, str) or not requested.strip(): + return None + requested = requested.strip() + try: + index = _index() + entry = index.get(requested.lower()) + if entry is not None: + variant = entry.variants[0] if entry.variants else None + return entry.load_path, variant, entry.loader_id + + base, sep, variant = requested.rpartition(":") + if not sep: + return None + entry = index.get(base.strip().lower()) + if entry is None: + return None + wanted = variant.strip().lower() + for v in entry.variants: + if v.lower() == wanted: + return entry.load_path, v, entry.loader_id + return None + except Exception: + # Best-effort: any resolver failure falls through to the loaded model, + # so a malformed name can never turn a servable request into a 500. + return None diff --git a/studio/backend/main.py b/studio/backend/main.py index 8b8a5fe787..5402e5eb7b 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -532,6 +532,11 @@ async def lifespan(app: FastAPI): _start_helper_precache_if_enabled() threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() + # Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set). + from core.inference.llama_keepwarm import idle_unload_loop + + app.state.idle_unload_task = asyncio.create_task(idle_unload_loop()) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair @@ -561,6 +566,14 @@ async def lifespan(app: FastAPI): ) yield + _idle_task = getattr(app.state, "idle_unload_task", None) + if _idle_task is not None: + _idle_task.cancel() + try: + await _idle_task + except asyncio.CancelledError: + pass + from core.inference.llama_http import aclose as _close_llama_http await _close_llama_http() @@ -883,6 +896,11 @@ app.add_middleware( upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes, ) +# Tracks in-flight inference requests for idle auto-unload; off -> passthrough. +from core.inference.llama_keepwarm import LlamaKeepWarmMiddleware # noqa: E402 + +app.add_middleware(LlamaKeepWarmMiddleware) + from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9caacec61b..9d9db83543 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -19,6 +19,7 @@ import httpx from loggers import get_logger import asyncio import threading +import weakref import re as _re @@ -2266,6 +2267,411 @@ def get_llama_cpp_backend() -> LlamaCppBackend: return _llama_cpp_backend +# Serializes opt-in auto-switch loads so two requests can't race a swap. One +# lock per running loop, since a module-level asyncio.Lock binds to a single +# loop and breaks multi-loop runners (e.g. pytest's per-test loops on pre-3.10). +_auto_switch_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_auto_switch_locks_guard = threading.Lock() + + +def _auto_switch_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + # WeakKeyDictionary mutation isn't thread-safe; guard get-or-create so two + # loops on different threads can't race it. + with _auto_switch_locks_guard: + lock = _auto_switch_locks.get(loop) + if lock is None: + lock = _auto_switch_locks[loop] = asyncio.Lock() + return lock + + +# Process-wide gate so a swap on another event loop in this process can't race +# this one for the single model slot: the asyncio lock above is per loop, but the +# backend slot and _load_model_impl are process-wide. threading.Lock so it serializes +# across loops/threads; released from the loop thread (Lock allows cross-thread release). +_auto_switch_process_lock = threading.Lock() + + +async def _acquire_swap_gate() -> None: + # Non-blocking first for the common single-loop case; otherwise poll off a + # short sleep rather than awaiting to_thread(acquire). A cancelled to_thread + # (client disconnect mid-wait) leaves its worker thread still acquiring, so the + # gate gets taken but the finally that releases it never runs -- deadlocking + # later swaps. Polling keeps the wait off this loop AND cancellation-safe: a + # cancel lands during the sleep, when the gate is not held. + while not _auto_switch_process_lock.acquire(blocking = False): + await asyncio.sleep(0.02) + + +# Counts in-flight auto-switch requests per (target, variant). The busy guard +# subtracts same-target waiters so concurrent requests for one model load once +# instead of each 409-ing the other. +_auto_switch_waiters: dict[tuple[str, str], int] = {} +_auto_switch_waiters_guard = threading.Lock() + + +def _switch_key(override_id: str, variant: Optional[str]) -> tuple[str, str]: + return (override_id.lower(), (variant or "").lower()) + + +def _note_switch_waiter(key: tuple[str, str], delta: int) -> None: + with _auto_switch_waiters_guard: + n = _auto_switch_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_waiters[key] = n + else: + _auto_switch_waiters.pop(key, None) + + +def _same_target_waiters(key: tuple[str, str]) -> int: + with _auto_switch_waiters_guard: + return _auto_switch_waiters.get(key, 0) + + +# A second waiter map keyed by the raw requested model, registered before the +# (slow) resolve. The middleware counts a concurrent same-model request as +# in-flight before it resolves and joins _auto_switch_waiters, so without this +# the first request would see it as an unrelated request and 409. +_auto_switch_request_waiters: dict[str, int] = {} +_auto_switch_request_waiters_guard = threading.Lock() + + +def _request_waiter_key(requested_model: str) -> str: + return requested_model.strip().lower() + + +def _note_request_waiter(key: str, delta: int) -> None: + with _auto_switch_request_waiters_guard: + n = _auto_switch_request_waiters.get(key, 0) + delta + if n > 0: + _auto_switch_request_waiters[key] = n + else: + _auto_switch_request_waiters.pop(key, None) + + +def _same_request_waiters(key: str) -> int: + with _auto_switch_request_waiters_guard: + return _auto_switch_request_waiters.get(key, 0) + + +def _llama_public_model_id(llama_backend, fallback: Optional[str] = None) -> Optional[str]: + """The id to report for the loaded GGUF in API responses: the advertised repo + id from an auto-switch load, else the cleaned public id, never the on-disk + .gguf path (see core.inference.model_ids.public_model_id).""" + return ( + getattr(llama_backend, "_openai_advertised_id", None) + or public_model_id(getattr(llama_backend, "model_identifier", None)) + or public_model_id(fallback) + or fallback + ) + + +_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY = "_unsloth_disable_openai_auto_switch" +# Sentinel a raw-body endpoint passes when the request omits ``model``: it must +# only restore an idle-freed model, never run the resolver (so a downloaded GGUF +# literally named "default" can't be swapped to). The NUL keeps it off any index. +_RELOAD_ONLY_MODEL = "\x00reload-only" + + +def _switch_model_for_payload(payload) -> str: + # A pydantic request fills an omitted ``model`` with "default"; only an + # explicitly set model may switch, else reload-only so a GGUF named "default" + # is never matched (mirrors the raw-body sentinel path). + return payload.model if "model" in payload.model_fields_set else _RELOAD_ONLY_MODEL + + +def _target_is_vision(load_path: str) -> bool: + # A local GGUF's vision capability is its companion mmproj, a filesystem check + # (no model load). Matches the loaded backend's is_vision, so rejecting a swap + # here can't differ from the post-load guard. Thread the ambient HF token so the + # probe keeps the capability-probe invariant (the resolver only yields local + # paths, where the token is unused, but the rule requires it regardless). + from utils.models.model_config import is_vision_model + try: + return bool(is_vision_model(load_path, hf_token = os.environ.get("HF_TOKEN"))) + except Exception as exc: + # Detection failure: don't block the swap, let the load decide. + logger.debug("auto-switch: vision probe failed for %s: %s", load_path, exc) + return True + + +def _messages_have_image(messages) -> bool: + return any( + isinstance(m.content, list) and any(isinstance(p, ImageContentPart) for p in m.content) + for m in messages + ) + + +def _request_has_image(payload) -> bool: + if getattr(payload, "image_base64", None): + return True + return _messages_have_image(payload.messages) + + +def _anthropic_request_has_image(payload) -> bool: + # Mirror anthropic_messages_to_openai: an Anthropic image block carries + # ``type == "image"`` (typed AnthropicImageBlock or a raw dict). + for msg in getattr(payload, "messages", None) or []: + content = getattr(msg, "content", None) + if not isinstance(content, list): + continue + for block in content: + bt = block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + if bt == "image": + return True + return False + + +def disable_openai_auto_switch_for_request(scope) -> None: + """Opt a request out of OpenAI auto-switch. The public preview route uses this: + it always serves its pinned checkpoint, so a caller-supplied model must never + swap the loaded model.""" + if isinstance(scope, dict): + scope[_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY] = True + + +def _automatic_model_load_may_run() -> bool: + """True when a request can trigger an automatic load: either resolver-based + auto-switch is on, or a standalone idle TTL can reload an idle-freed model. The + validate-before-switch guards key off this so an invalid request never loads.""" + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + ) + return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 + + +async def _maybe_auto_switch_model( + requested_model: Optional[str], + fastapi_request: Request, + current_subject: str, + *, + require_vision: bool = False, +) -> None: + """Load a downloaded local GGUF named by an OpenAI request when auto-switch is on. + + No-op unless enabled and ``requested_model`` resolves to a downloaded local + model different from the loaded one. Unknown names fall through (drop-in + compat) and no remote download is triggered. ``require_vision`` rejects a swap + to a text-only target before it runs, so an image request can't evict the + resident vision model only to 400 afterwards. + """ + from utils.openai_auto_switch_settings import ( + get_openai_auto_switch_enabled, + get_auto_unload_idle_seconds, + get_model_override, + ) + from core.inference.local_model_resolver import resolve_local_gguf + from core.inference.llama_keepwarm import ( + get_last_unloaded_model, + other_inference_request_count, + inference_lifecycle_gate, + ) + + # Treat a non-string model (e.g. {"model": 123} on a raw-body endpoint) as + # absent so it falls through instead of raising in the membership checks below. + if not isinstance(requested_model, str) or not requested_model: + return + # The public preview route opts out so a caller cannot switch away from the + # pinned preview checkpoint it just loaded. + scope = getattr(fastapi_request, "scope", None) + if isinstance(scope, dict) and scope.get(_DISABLE_OPENAI_AUTO_SWITCH_SCOPE_KEY): + return + auto_switch_on = get_openai_auto_switch_enabled() + # The reload-stash path also runs when idle-unload is active on its own (a + # standalone UNSLOTH_MODEL_IDLE_TTL with auto-switch off), so a model the idle + # loop freed is restored on the next request. The resolver-based switch still + # requires the auto-switch toggle. + if not auto_switch_on and get_auto_unload_idle_seconds() <= 0: + return + + # Register by the raw requested model before resolving (which can be slow): + # the middleware already counts a concurrent same-model request as in-flight, + # so the busy guard must know it shares this target even while it resolves. + request_key = _request_waiter_key(requested_model) + _note_request_waiter(request_key, 1) + try: + # Off the loop: a cold-cache rebuild walks several model dirs + HF caches. + # With auto-switch off (or an omitted-model reload-only request), skip the + # resolve so only the reload-stash path runs and no name is ever matched. + reload_only = requested_model == _RELOAD_ONLY_MODEL + resolved = ( + await asyncio.to_thread(resolve_local_gguf, requested_model) + if auto_switch_on and not reload_only + else None + ) + if resolved is None: + # Idle-unload may have freed the model; reload exactly what it freed + # (path + quant + advertised id) so an alias/unknown name stays servable + # and keeps the override keyed by the advertised id, not the load path. + last = get_last_unloaded_model() + # A non-GGUF (Unsloth/Transformers) model loaded after the idle-unload + # leaves the GGUF slot empty but is the live model, so don't resurrect + # the stale GGUF over it (that load would tear the active model down). + if ( + not last + or get_llama_cpp_backend().is_loaded + or getattr(get_inference_backend(), "active_model_name", None) + ): + return + if len(last) == 3: + target_id, variant, override_id = last + else: # pre-3-tuple stash: fall back to the path as the override key + target_id, variant = last + override_id = target_id + else: + # load_path is a concrete local path (never the bare repo id), so /load + # takes the local branch and cannot trigger a download. override_id is the + # advertised repo id, the launch-override key and the public model id. + target_id, variant, override_id = resolved + backend = get_llama_cpp_backend() + # A bare model id (no :VARIANT) is satisfied by any loaded quant of that + # repo, so it never reloads a different local quant that already serves it. + bare = ":" not in requested_model + + def _already_serving() -> bool: + # Match against both the concrete load path and the advertised repo id, + # so a model loaded manually by repo id (identifier = repo id) and one + # loaded by auto-switch (identifier = path, advertised = repo id) both + # count as already serving rather than triggering a needless reswap. + if not backend.is_loaded or not backend.model_identifier: + return False + loaded_keys = {backend.model_identifier.lower()} + advertised = getattr(backend, "_openai_advertised_id", None) + if advertised: + loaded_keys.add(advertised.lower()) + if loaded_keys.isdisjoint({target_id.lower(), override_id.lower()}): + return False + if bare: + return True + if variant: + loaded_variant = (getattr(backend, "hf_variant", None) or "").lower() + return loaded_variant == variant.lower() + return True + + def _record_serving_alias() -> None: + # When an advertised alias already resolves to the loaded model (e.g. a + # model loaded by local path, requested by its repo/LM Studio id), record + # the alias as the public id so /v1/models and responses report it (and + # mark it loaded) instead of the path-derived basename. Resolver branch + # only: the reload-stash override_id can be the bare path, not a repo id. + # Lock-free is safe here: an in-flight request blocks any concurrent swap + # (single-slot busy guard), so the loaded model can't change under this. + if resolved is None or not override_id: + return + b = get_llama_cpp_backend() + if getattr(b, "_openai_advertised_id", None) != override_id: + b._openai_advertised_id = override_id + + if _already_serving(): + _record_serving_alias() + return + # An image/audio request naming a different text-only GGUF would load it + # here and only 400 below, evicting the working model. Reject before the + # swap. Only the resolver branch (an explicit new target); the reload-stash + # path just restores the model the request was already using. Both vision and + # audio input come from a companion mmproj (a filesystem probe) -- run it off + # the loop, like the resolver above. + if ( + require_vision + and resolved is not None + and not await asyncio.to_thread(_target_is_vision, target_id) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "The requested model does not support the image or audio input in this request.", + status = 400, + code = "invalid_value", + param = "model", + ), + ) + key = _switch_key(override_id, variant) + _note_switch_waiter(key, 1) + try: + async with _auto_switch_lock(): + # The asyncio lock is per loop; add a process-wide gate so a swap on + # another loop in this process can't race the single slot. + await _acquire_swap_gate() + try: + # Hold the keep-warm gate across the swap so no new inference can + # start on the model while it is being torn down and replaced. + async with inference_lifecycle_gate(): + if _already_serving(): + _record_serving_alias() + return + # Single slot: refuse a cross-model swap while another inference + # request is active rather than killing its response. Requests + # heading to this same target (by resolved id or raw name) are + # excluded, so concurrent requests for one model load once. A + # pending request is still in the middleware, not generating, so + # it is not counted here. + same_others = max( + _same_target_waiters(key) - 1, _same_request_waiters(request_key) - 1, 0 + ) + others = other_inference_request_count( + current_request_counted = True, include_pending = False + ) + # Not gated on the GGUF being loaded: _load_model_impl also + # tears down an active Unsloth backend before loading a GGUF, + # so refuse whenever any other inference request is in flight. + if others > same_others: + raise HTTPException( + status_code = 409, + detail = openai_error_body( + "Cannot switch models while another inference request is in progress.", + status = 409, + code = "model_switch_busy", + param = "model", + ), + ) + # Apply this model's saved launch flags so the swap honors the config. + override = get_model_override(override_id) + load_kwargs = {"model_path": target_id, "gguf_variant": variant} + if override.get("llama_extra_args") is not None: + load_kwargs["llama_extra_args"] = override["llama_extra_args"] + if override.get("max_seq_length") is not None: + load_kwargs["max_seq_length"] = override["max_seq_length"] + # Reuse the load impl so its dedup, tensor fallback, and threading + # apply. Call the impl directly: we already hold the lifecycle gate + # the /load route would otherwise take, so the route would deadlock. + await _load_model_impl( + LoadRequest(**load_kwargs), + fastapi_request, + current_subject, + ) + # Advertise the repo id (not the concrete load path) as the loaded + # model's public id and override key for /v1/models and idle stash. + get_llama_cpp_backend()._openai_advertised_id = override_id + finally: + _auto_switch_process_lock.release() + finally: + _note_switch_waiter(key, -1) + finally: + _note_request_waiter(request_key, -1) + + +async def _auto_switch_from_request_body(request: Request, current_subject: str): + """Run auto-switch from a raw-body endpoint's ``model`` without changing its + pre-feature status codes: a malformed/non-dict body yields no model (so an + unloaded backend still 503s, not 500), and the caller re-reads to surface the + original parse error after the loaded-state check. Returns the parsed body, or + None if it could not be parsed.""" + try: + body = await request.json() + except (json.JSONDecodeError, ValueError): + return None + if isinstance(body, dict): + # A raw-body client may omit ``model`` and rely on the loaded backend. Pass + # a reload-only sentinel so the idle-stash reload still runs (an idle-freed + # model is restored) without the resolver ever matching a real name. + model = body.get("model") or _RELOAD_ONLY_MODEL + else: + model = None + await _maybe_auto_switch_model(model, request, current_subject) + return body + + def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool: """Effective quantization the loader will use: a LoRA adapter can flip 4-bit to 16-bit via adapter_config.json, so the guard sizes this, not the raw request.""" @@ -2508,6 +2914,15 @@ async def load_model( GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ + # Hold the lifecycle gate across the load so idle auto-unload can't unload the + # model mid-load. Auto-switch calls _load_model_impl directly since it already + # holds this gate. + from core.inference.llama_keepwarm import inference_lifecycle_gate + async with inference_lifecycle_gate(): + return await _load_model_impl(request, fastapi_request, current_subject) + + +async def _load_model_impl(request: LoadRequest, fastapi_request: Request, current_subject: str): from core.inference.llama_cpp import LlamaServerNotFoundError native_grant_backed = False @@ -2932,6 +3347,13 @@ async def load_model( logger.info( f"Loaded GGUF model via llama-server: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash now, not only on the next poll. + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() + # A plain load advertises its own identifier; auto-switch overwrites + # this with the repo id right after _load_model_impl returns. + llama_backend._openai_advertised_id = None # Audio detection moved into load_model under _serial_load_lock (#5642). _gguf_audio = llama_backend._audio_type @@ -3033,6 +3455,13 @@ async def load_model( logger.info( f"Loaded model: {model_log_label if native_grant_backed else config.identifier}" ) + # Clear any idle-unload reload stash: a manual load supersedes an idle-freed + # GGUF, so the next /v1 request must not resurrect it. Mirror the GGUF branch + # above; without this a non-GGUF load leaves a stale stash until the idle + # poll clears it (and never, while idle-unload is off). + from core.inference.llama_keepwarm import note_model_loaded + + note_model_loaded() # Load inference configuration parameters inference_config = load_inference_config(config.identifier) @@ -3363,6 +3792,10 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge Unload a model from memory. Routes to the correct backend (llama-server for GGUF, Unsloth otherwise). """ + # A deliberate unload means "stay unloaded": drop any idle reload stash so the + # next /v1 request can't resurrect this model. The idle loop unloads via the + # backend directly (not this route), so clearing here never fights keep-warm. + from core.inference.llama_keepwarm import note_model_unloaded try: # Check if the GGUF backend has this model loaded or is loading it. llama_backend = get_llama_cpp_backend() @@ -3371,13 +3804,17 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) or not llama_backend.is_loaded ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). llama_backend.unload_model() + note_model_unloaded() logger.info(f"Unloaded GGUF model: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) # Otherwise, unload from Unsloth backend backend = get_inference_backend() backend.unload_model(request.model_path) + note_model_unloaded() logger.info(f"Unloaded model: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) @@ -3786,10 +4223,25 @@ async def generate_audio( raise HTTPException(status_code = 400, detail = "No user message found.") text = last_user_msg["content"] + # Restore an idle-evicted GGUF before selecting a backend: this path is + # keep-warm-tracked but had no reload hook, so a standalone idle TTL could + # unload an audio GGUF the next request then failed to restore. Validation + # above ran first, so an invalid request never triggers a reload. + # + # Reload-only on purpose: a local GGUF's audio-input capability is not a cheap + # pre-load probe (the companion mmproj signal can't tell an audio projector + # from a vision one, and codec-based TTS ships no projector at all), so passing + # the client model through the resolver could load a text- or vision-only target + # and evict the working audio model before the audio backend check fails. Only + # the idle-stash restore runs here; switching TTS models is an explicit /load. + await _maybe_auto_switch_model(_RELOAD_ONLY_MODEL, request, current_subject) + # Pick backend — both return (wav_bytes, sample_rate) llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False): - model_name = public_model_id(llama_backend.model_identifier) + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend) gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -4815,6 +5267,12 @@ async def openai_chat_completions( # ── External provider routing ──────────────────────────────── # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth. if payload.provider_id or payload.provider_type: + # External provider: this request won't touch the local GGUF, so drop it + # from the keep-warm count or its in-flight stream would falsely block a + # concurrent local auto-switch with model_switch_busy. + from core.inference.llama_keepwarm import untrack_current_request + + untrack_current_request(request.scope) # Bypass Permissions suppresses the confirm gate, so do not reject a # request that sets both flags (effective confirm is then False). if ( @@ -4868,6 +5326,95 @@ async def openai_chat_completions( ), ) + # Reject a system-only chat before any automatic load so an invalid request + # never swaps or reloads the resident model (as /responses and /messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + # Parse once and reuse below. + _pre_parsed = None + _needs_vision = False + if _automatic_model_load_may_run(): + _pre_parsed = _extract_content_parts(payload.messages) + if not _pre_parsed[1]: + raise HTTPException( + status_code = 400, detail = "At least one non-system message is required." + ) + # Reject confirm-without-stream local tool requests before the switch: the + # local tool path requires stream=true for the confirm gate, so this shape + # is invalid and must not evict the resident model first. Mirror that path's + # enablement exactly (_effective_enable_tools honors a CLI --enable-tools + # policy hard-override; mcp_enabled opens the tool loop on its own but still + # defers to a CLI --disable-tools policy), or an mcp_enabled/policy-forced + # request would slip past this guard and only 400 after the swap. + from state.tool_policy import get_tool_policy as _get_confirm_tool_policy + + _confirm_cli_policy = _get_confirm_tool_policy() + if ( + payload.confirm_tool_calls + and not payload.bypass_permissions + and not payload.stream + and ( + _effective_enable_tools(payload) + or (bool(payload.mcp_enabled) and _confirm_cli_policy is not False) + or bool(payload.enabled_tools) + or bool(payload.tools) + or bool(payload.openai_code_exec_container_id) + or bool(payload.anthropic_code_exec_container_id) + ) + ): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "confirm_tool_calls requires stream=true for local tool execution.", + status = 400, + code = "invalid_request_error", + param = "confirm_tool_calls", + ), + ) + # Reject a malformed tool_choice forcing object before the switch: a + # {"type": "function", "function": {}} with no name would otherwise be + # forwarded to llama-server and rejected only after the model swapped. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_fn = _tc.get("function") + _tc_name = _tc_fn.get("name") if isinstance(_tc_fn, dict) else None + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # Reject an oversized audio upload before the switch: the size cap is a + # cheap, target-independent length check, so a too-large payload must not + # load a GGUF only to 413 afterward (the decode itself stays post-switch to + # avoid decoding a valid upload twice). + if payload.audio_base64 and len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException(status_code = 413, detail = "Audio file is too large (max ~25 MB).") + # Reject streaming n>1 before the switch: only the non-streaming GGUF path + # returns multiple choices, so stream=true + n>1 is invalid on every local + # serving path (the external path already rejected it before its early + # return). Both fields are known here, so a bad shape must not load model B + # only to 400. The non-streaming n>1 cases stay post-switch, where the + # serving path decides whether the shape is supported. + if payload.stream and _wants_multiple_choices(payload): + _raise_unsupported_n("streaming chat completions") + # Audio input rides the same companion-mmproj projector as vision, so a + # text-only target can't serve it either; guard both before the switch. + _needs_vision = ( + bool(_pre_parsed[2]) or _request_has_image(payload) or bool(payload.audio_base64) + ) + + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _needs_vision, + ) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded @@ -4923,8 +5470,9 @@ async def openai_chat_completions( return response if using_gguf: - # Echo a clean public id in the response, never the absolute .gguf path. - model_name = public_model_id(llama_backend.model_identifier) or payload.model + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path. + model_name = _llama_public_model_id(llama_backend, payload.model) if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") @@ -5185,7 +5733,11 @@ async def openai_chat_completions( ) # ── Parse messages (handles multimodal content parts) ───── - system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) + # Reuse the pre-hook parse when auto-switch did it, else parse now. + if _pre_parsed is not None: + system_prompt, chat_messages, extracted_image_b64 = _pre_parsed + else: + system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: raise _reject(400, "At least one non-system message is required.") @@ -6492,10 +7044,13 @@ def _openai_model_objects() -> list[dict]: # Check GGUF backend llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: + # Advertise the repo id an auto-switch load recorded, not the concrete + # on-disk load path, so /v1/models never leaks a host path or lists a + # model twice (path plus repo id). entry = { - # Public id, never the absolute .gguf path (which leaks the host - # filesystem layout); see core.inference.model_ids.public_model_id. - "id": public_model_id(llama_backend.model_identifier), + # Advertised repo id after an auto-switch load, else a clean public id, + # never the absolute .gguf path (which leaks the host filesystem layout). + "id": _llama_public_model_id(llama_backend), "object": "model", "created": _created, "owned_by": _OWNED_BY, @@ -6541,7 +7096,21 @@ def _openai_model_objects() -> list[dict]: # don't rescan the HF cache and models dirs on every request. _CATALOG_CACHE: dict = {"at": 0.0, "models": []} _CATALOG_TTL_S = 30.0 -_CATALOG_LOCK = asyncio.Lock() +# Per-loop lock (like _auto_switch_lock): a module-level asyncio.Lock ties its +# waiters to the loop that first awaited it, so a second event loop awaiting it +# in a multi-loop ASGI process can hang. The cache double-check keeps correctness +# even when two loops each scan once. +_catalog_locks: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_catalog_locks_guard = threading.Lock() + + +def _catalog_lock() -> asyncio.Lock: + loop = asyncio.get_running_loop() + with _catalog_locks_guard: + lock = _catalog_locks.get(loop) + if lock is None: + lock = _catalog_locks[loop] = asyncio.Lock() + return lock async def _cached_local_catalog() -> list: @@ -6558,7 +7127,7 @@ async def _cached_local_catalog() -> list: now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] - async with _CATALOG_LOCK: + async with _catalog_lock(): now = time.monotonic() if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: return _CATALOG_CACHE["models"] @@ -6588,7 +7157,15 @@ async def _openai_catalog_objects() -> list[dict]: by_id[entry["id"]] = {**entry, "loaded": True} # Locally available (downloaded/cached) models that are not already loaded. - for info in await _cached_local_catalog(): + # Advertise only GGUF models /v1 can actually serve (llama.cpp). GGUF-ness is + # read from the on-disk files, not model_format: the HF-cache scanner leaves + # model_format unset for GGUF snapshots, so a model_format filter would drop + # every cached GGUF. The file checks run off the loop. + from core.inference.local_model_resolver import info_has_local_gguf + + catalog = await _cached_local_catalog() + servable = await asyncio.to_thread(lambda: [i for i in catalog if info_has_local_gguf(i)]) + for info in servable: cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None)) if not cid or cid in by_id: continue @@ -6632,29 +7209,43 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge from core.inference.model_ids import model_id_matches # Loaded models resolve without a catalog scan (the common case); only build - # the full catalog -- which may hit the filesystem -- for unloaded ids. - for entry in _openai_model_objects(): - if entry["id"] == model_id: + # the full catalog -- which may hit the filesystem -- for unloaded ids. Match + # case-insensitively, like the catalog loop below and the resolver's index. + _loaded = _openai_model_objects() + for entry in _loaded: + eid = entry["id"] + if isinstance(eid, str) and eid.lower() == model_id.lower(): return {**entry, "loaded": True} objects = await _openai_catalog_objects() for model in objects: - if model["id"] == model_id: + # Case-insensitive to match the resolver, which lowercases its index. + mid = model.get("id") + if isinstance(mid, str) and mid.lower() == model_id.lower(): return model # Backward compatibility: a client may still send the legacy raw identifier - # (e.g. an absolute .gguf path cached from an older /v1/models). Resolve it to - # the clean object so it keeps working, without ever echoing the path back. + # (e.g. an absolute .gguf path cached from an older /v1/models). Map it to the + # loaded model's object so it keeps working, without ever echoing the path back. + # Key each raw id to the SAME public id its /v1/models entry uses: an + # auto-switch load advertises a repo id while its identifier is the snapshot + # path, so public_model_id(path) would miss the advertised entry and 404 a + # model that is in fact loaded. llama_backend = get_llama_cpp_backend() backend = get_inference_backend() - for raw in ( - llama_backend.model_identifier if llama_backend.is_loaded else None, - backend.active_model_name or None, - ): - if raw and model_id_matches(model_id, raw): - clean = public_model_id(raw) - for model in objects: - if model["id"] == clean: - return model + raw_to_public: list[tuple[str, Optional[str]]] = [] + if llama_backend.is_loaded and llama_backend.model_identifier: + raw_to_public.append( + (llama_backend.model_identifier, _llama_public_model_id(llama_backend)) + ) + if backend.active_model_name: + raw_to_public.append( + (backend.active_model_name, public_model_id(backend.active_model_name)) + ) + for raw, clean in raw_to_public: + if model_id_matches(model_id, raw): + for entry in _loaded: + if entry["id"] == clean: + return {**entry, "loaded": True} raise HTTPException( status_code = 404, detail = openai_error_body( @@ -6679,6 +7270,16 @@ def _flatten_monitor_prompt(value) -> str: return str(value) +def _completions_prompt_present(body: dict) -> bool: + """Whether a completions body carries a usable ``prompt`` (non-empty).""" + prompt = body.get("prompt") + if isinstance(prompt, str): + return prompt != "" + if isinstance(prompt, (list, tuple)): + return len(prompt) > 0 + return prompt is not None + + @router.post("/completions") async def openai_completions(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6688,13 +7289,39 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge when a GGUF model is loaded. """ llama_backend = get_llama_cpp_backend() + + # Reject a request with no prompt before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/embeddings already + # validate before switching). Gate on every automatic-load trigger. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_prompt = _pre.get("prompt") + if _pre_prompt is not None and not isinstance(_pre_prompt, (str, list, tuple)): + # An object/number prompt is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'prompt' must be a string or array.") + if not _completions_prompt_present(_pre): + raise HTTPException(status_code = 400, detail = "'prompt' is required for completions.") + + # Opt-in: load the requested local GGUF before the loaded-state check. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() if body.get("max_tokens") is None: body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" @@ -6703,7 +7330,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -6850,6 +7477,16 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge # ===================================================================== +def _embeddings_input_present(body: dict) -> bool: + """Whether an embeddings body carries a usable ``input`` (non-empty).""" + inp = body.get("input") + if isinstance(inp, str): + return inp != "" + if isinstance(inp, (list, tuple)): + return len(inp) > 0 + return inp is not None + + @router.post("/embeddings") async def openai_embeddings(request: Request, current_subject: str = Depends(get_current_subject)): """ @@ -6861,13 +7498,42 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get error (expected). """ llama_backend = get_llama_cpp_backend() + # Reject a request with no input before any automatic load so an invalid + # request never swaps or reloads the resident model (as chat/responses/messages + # already validate before switching). Gate on every automatic-load trigger, + # not just auto-switch, since a standalone idle TTL can also reload here. + if _automatic_model_load_may_run(): + try: + _pre = await request.json() + except (json.JSONDecodeError, ValueError): + _pre = None + if isinstance(_pre, dict): + _pre_input = _pre.get("input") + if _pre_input is not None and not isinstance(_pre_input, (str, list, tuple)): + # An object/number input is a deterministic client error (only a + # string or array is valid); reject it before the switch so a bad + # shape can't load a GGUF only to be rejected by llama-server after. + raise HTTPException(status_code = 400, detail = "'input' must be a string or array.") + if not _embeddings_input_present(_pre): + raise HTTPException(status_code = 400, detail = "'input' is required for embeddings.") + # Embeddings is a model-bearing inference path too, so honor auto-switch. Unlike + # vision (cheaply pre-checked via a companion mmproj), GGUF pooling capability has + # no reliable pre-load probe -- is_embedding_model keys on a sentence-transformers + # modules.json a bare .gguf never has -- so embeddings auto-switch is best-effort: + # a non-embedding target switches, then llama-server returns a no-pooling error. + body = await _auto_switch_from_request_body(request, current_subject) if not llama_backend.is_loaded: raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) + if not isinstance(body, dict): + # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); + # a valid non-dict body such as a list is a clean 400 rather than a 500. + body = await request.json() + if not isinstance(body, dict): + raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - body = await request.json() target_url = f"{llama_backend.base_url}/v1/embeddings" prompt_text = _flatten_monitor_prompt(body.get("input", "")) monitor_id = None @@ -6875,7 +7541,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get monitor_id = api_monitor.start( endpoint = request.url.path, method = request.method, - model = str(body.get("model") or llama_backend.model_identifier or "default"), + model = str(body.get("model") or _llama_public_model_id(llama_backend) or "default"), prompt = prompt_text, context_length = llama_backend.context_length, subject = current_subject, @@ -7340,10 +8006,13 @@ def _build_chat_request( ``/v1/chat/completions`` client-side pass-through picks them up unchanged. """ chat_kwargs: dict = dict( - model = payload.model, messages = messages, stream = stream, ) + # Only forward an explicitly set model so an omitted Responses model stays + # reload-only when openai_chat_completions re-checks on the non-streaming path. + if "model" in payload.model_fields_set: + chat_kwargs["model"] = payload.model if payload.temperature is not None: chat_kwargs["temperature"] = payload.temperature if payload.top_p is not None: @@ -7597,12 +8266,11 @@ async def _responses_stream( # Clean public id for every response envelope. Prefer the loaded model's # id so the stream agrees with /v1/models, chat/completions and the # non-streaming twin; fall back to a sanitized payload.model (a legacy - # raw .gguf path is stripped, never echoed back). - _clean_model = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model - ) + # raw .gguf path is stripped, never echoed back). Use the advertised-id + # helper, not the raw identifier: after an auto-switch to a cached HF GGUF + # the identifier is the snapshot path while the repo id lives in + # _openai_advertised_id, so the raw form would stream a snapshot basename. + _clean_model = _llama_public_model_id(llama_backend, payload.model) or payload.model full_text = "" full_reasoning = "" input_tokens = 0 @@ -8245,6 +8913,56 @@ async def openai_responses( messages = _normalise_responses_input(payload) if not messages: raise HTTPException(status_code = 400, detail = "No input provided.") + # System/developer-only input normalises to a non-empty list, so reject it + # before the switch (mirror chat) or an invalid request evicts the resident + # model only for the chat handler to 400 it as having no non-system message. + if not any(m.role not in ("system", "developer") for m in messages): + raise HTTPException(status_code = 400, detail = "At least one non-system message is required.") + # Reject a malformed function tool before any model load, mirroring the + # /v1/chat/completions check, so an invalid request never switches the model. + # Built-in tools (web_search, mcp, ...) carry no name and are dropped later. + for _tool in payload.tools or []: + if not isinstance(_tool, dict) or _tool.get("type") != "function": + continue + _name = _tool.get("name") + if not isinstance(_name, str) or not _name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tools': each function tool must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tools", + ), + ) + # Reject a forcing-function tool_choice with no name before the switch (mirror + # chat), so a malformed request can't evict the model. Responses forces with + # {"type": "function", "name": "X"}; the streaming path would otherwise forward + # the bad choice and the non-streaming path only 400s after the swap. + _tc = payload.tool_choice + if isinstance(_tc, dict) and _tc.get("type") == "function": + _tc_name = _tc.get("name") + if not isinstance(_tc_name, str) or not _tc_name.strip(): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Invalid 'tool_choice': the forced function must have a 'name'.", + status = 400, + code = "invalid_value", + param = "tool_choice", + ), + ) + # After input validation so a 400 never triggers a load. Switches the + # streaming path; non-streaming re-checks via the idempotent chat handler. + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to 400 afterwards + # (the non-streaming chat re-check short-circuits on _already_serving). + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _messages_have_image(messages), + ) if payload.stream: monitor_id = None @@ -8381,6 +9099,28 @@ def _normalize_anthropic_openai_images(openai_messages: list[dict], is_vision: b return has_image +def _validate_anthropic_client_tools(tools) -> None: + # Reject malformed client tools before any model load, so an invalid request + # never evicts the loaded model. AnthropicTool relaxed name/input_schema to + # Optional for server tools, so the converter silently drops incomplete + # entries; surface them as 400 here. A `type` field marks a server-tool + # declaration (unrecognized server tools are no-ops); anything else without + # input_schema or name is malformed. + for tool in tools or []: + td = tool if isinstance(tool, dict) else tool.model_dump() + name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") + if schema is None and not isinstance(type_, str): + raise HTTPException( + status_code = 400, + detail = f"Tool {name!r} is missing required field 'input_schema'.", + ) + if schema is not None and (not isinstance(name, str) or not name): + raise HTTPException( + status_code = 400, + detail = "Client tool is missing required field 'name'.", + ) + + @router.post("/messages/count_tokens") async def anthropic_count_tokens( payload: AnthropicMessagesRequest, @@ -8394,6 +9134,19 @@ async def anthropic_count_tokens( tokenizer, and returns ``{"input_tokens": int}`` only. Unlike /messages, max_tokens is NOT required here. """ + # Reject malformed tools before the switch, like /messages, so an invalid + # count request can't evict the loaded model. + _validate_anthropic_client_tools(payload.tools) + # Count with the requested model's tokenizer, like the sibling /messages. + # Carry the vision guard too: an image count naming a text-only GGUF must not + # evict a loaded vision model for a swap that can't serve the request. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + llama_backend = get_llama_cpp_backend() if not llama_backend.is_loaded: raise HTTPException( @@ -8460,14 +9213,20 @@ async def anthropic_messages( JSON). """ llama_backend = get_llama_cpp_backend() - if not llama_backend.is_loaded: + + # Default-off parity: with no automatic load possible and nothing loaded, 503 + # before any request-shape check, exactly as the pre-feature endpoint did. When + # an automatic load can run (auto-switch or a standalone idle TTL), fall through + # so validation runs before the reload hook gets a chance to restore the model. + if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, detail = "No GGUF model loaded. Load a GGUF model first.", ) - # max_tokens is a required field on the Anthropic Messages API; real - # Anthropic returns a 400 invalid_request_error when it is omitted. + # max_tokens is a required field on the Anthropic Messages API; real Anthropic + # returns a 400 invalid_request_error when it is omitted. Validate before + # auto-switch so a rejected request never triggers a model load. if payload.max_tokens is None: raise HTTPException( status_code = 400, @@ -8478,13 +9237,47 @@ async def anthropic_messages( ), ) - # Clean public id so /v1/messages never echoes the local .gguf path (and a - # legacy raw path sent as payload.model is sanitized rather than returned). - model_name = ( - public_model_id(getattr(llama_backend, "model_identifier", None)) - or public_model_id(payload.model) - or payload.model + # Reject malformed client tools before any model load (see helper), so an + # invalid request never evicts the loaded model. + _validate_anthropic_client_tools(payload.tools) + + # Mixing Anthropic server tools with custom client tools is unsupported (the + # server-tool loop can't relay client functions back to the caller). Reject + # before the switch too -- it depends only on the payload -- so an invalid + # request never evicts the loaded model. Reused below for tool routing. + requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) + _has_client_tool = any( + (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None + for t in payload.tools or [] ) + if requested_studio_tools and _has_client_tool: + raise HTTPException( + status_code = 400, + detail = ( + "Mixing Anthropic server tools (e.g. web_search_20250305) " + "with custom client tools in a single request is not " + "supported. Send them in separate requests." + ), + ) + + # require_vision rejects a swap to a text-only target before it runs, so an + # image request can't evict the resident vision model only to hit the vision + # guard (_normalize_anthropic_openai_images) below after the load. + await _maybe_auto_switch_model( + _switch_model_for_payload(payload), + request, + current_subject, + require_vision = _anthropic_request_has_image(payload), + ) + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + # Advertised repo id after an auto-switch load, else a clean public id, never + # the local .gguf path (and a legacy raw path in payload.model is sanitized). + model_name = _llama_public_model_id(llama_backend, payload.model) message_id = f"msg_{uuid.uuid4().hex[:24]}" # ── Translate Anthropic → OpenAI ────────────────────────── @@ -8531,51 +9324,8 @@ async def anthropic_messages( # 2. tools=[...] only → client-side pass-through (standard Anthropic behavior) # 3. neither → plain chat # The server-side agentic loop doesn't support multimodal input -- matches - # the `not image_b64` gate in /v1/chat/completions. - requested_studio_tools = _anthropic_requested_studio_tools(payload.tools) - - # Reject malformed client tools at the boundary. AnthropicTool was relaxed - # to Optional[name]/Optional[input_schema] for server tools, so the - # converter silently drops incomplete entries -- surface them as 400. A - # `type` field marks a server-tool declaration per spec (unrecognized server - # tools are accepted as no-ops); anything else without input_schema or name - # is malformed and must not be allowed to silently flip execution mode or - # disable tool calling. - for tool in payload.tools or []: - td = tool if isinstance(tool, dict) else tool.model_dump() - name, type_, schema = td.get("name"), td.get("type"), td.get("input_schema") - if schema is None and not isinstance(type_, str): - raise HTTPException( - status_code = 400, - detail = f"Tool {name!r} is missing required field 'input_schema'.", - ) - if schema is not None and (not isinstance(name, str) or not name): - raise HTTPException( - status_code = 400, - detail = "Client tool is missing required field 'name'.", - ) - - # Detect client tools from the raw payload (presence of input_schema) so the - # mixed-mode check below isn't fooled by a name collision with a server-tool - # alias that the post-filter would silently drop. - _has_client_tool = any( - (t if isinstance(t, dict) else t.model_dump()).get("input_schema") is not None - for t in payload.tools or [] - ) - - # The server-tool agentic loop executes tools in-process and can't relay - # unknown client functions back to the caller, so mixed requests would - # silently drop the client tools. Reject explicitly instead. - if requested_studio_tools and _has_client_tool: - raise HTTPException( - status_code = 400, - detail = ( - "Mixing Anthropic server tools (e.g. web_search_20250305) " - "with custom client tools in a single request is not " - "supported. Send them in separate requests." - ), - ) - + # the `not image_b64` gate in /v1/chat/completions. requested_studio_tools and + # the mixed-mode rejection were computed before the switch above. openai_client_tools = [ tool for tool in anthropic_tools_to_openai(payload.tools or []) diff --git a/studio/backend/routes/preview.py b/studio/backend/routes/preview.py index 1fa9055d3a..5acf039401 100644 --- a/studio/backend/routes/preview.py +++ b/studio/backend/routes/preview.py @@ -17,7 +17,11 @@ from loggers import get_logger from auth.authentication import get_current_subject from auth.storage import DEFAULT_ADMIN_USERNAME from models.inference import ChatCompletionRequest, LoadRequest -from routes.inference import load_model, openai_chat_completions +from routes.inference import ( + disable_openai_auto_switch_for_request, + load_model, + openai_chat_completions, +) from state.tool_policy import tools_force_disabled from utils.client_ip import client_ip from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint @@ -155,6 +159,9 @@ async def _serve_chat( path = _resolve_or_4xx(run, checkpoint) is_lora = (path / "adapter_config.json").exists() payload = _sanitize_preview_payload(payload, is_lora) + # Preview always serves the pinned checkpoint it loads below; a public caller's + # `model` field must never trigger an OpenAI auto-switch to another GGUF. + disable_openai_auto_switch_for_request(getattr(request, "scope", None)) await _preview_lock.acquire() keep_locked = False try: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index f582500f7f..0694ae31e0 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,16 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.openai_auto_switch_settings import ( + DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, + DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, + get_auto_unload_idle_seconds, + get_model_overrides, + get_openai_auto_switch_enabled, + get_stored_auto_unload_idle_seconds, + set_model_override, + set_openai_auto_switch, +) from utils.preview_sharing_settings import ( DEFAULT_PREVIEW_SHARING_ENABLED, get_preview_sharing_enabled, @@ -66,6 +76,33 @@ class HelperPrecacheResponse(BaseModel): disabled_by_env: bool +class OpenAIAutoSwitchPayload(BaseModel): + enabled: bool + auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0) + + +class OpenAIAutoSwitchResponse(BaseModel): + enabled: bool + auto_unload_idle_seconds: int + default_enabled: bool = DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + # True when the idle-unload loop will actually unload (effective TTL > 0). With + # UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled + # is false, so the UI can show idle-unload as active instead of "needs enable". + idle_unload_active: bool = False + + +class ModelOverridePayload(BaseModel): + model_id: str = Field(..., min_length = 1) + llama_extra_args: list[str] = Field(default_factory = list) + # ge=1: 0 is not a valid sequence length, and the setter drops a falsy value, + # so reject it at the boundary instead of accepting then silently discarding it. + max_seq_length: Optional[int] = Field(default = None, ge = 1, le = 1048576) + + +class ModelOverridesResponse(BaseModel): + overrides: dict[str, dict] + + def _upload_limit_response(limit_mb: int) -> UploadLimitResponse: return UploadLimitResponse( max_upload_size_mb = limit_mb, @@ -128,6 +165,70 @@ def update_helper_precache( return _helper_precache_response(enabled) +@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def get_openai_auto_switch( + current_subject: str = Depends(get_current_subject), +) -> OpenAIAutoSwitchResponse: + return OpenAIAutoSwitchResponse( + enabled = get_openai_auto_switch_enabled(), + auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(), + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) +def update_openai_auto_switch( + payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject) +) -> OpenAIAutoSwitchResponse: + try: + enabled, idle_seconds = set_openai_auto_switch( + payload.enabled, payload.auto_unload_idle_seconds + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid OpenAI auto-switch setting."), + event = "settings.update_openai_auto_switch_failed", + log = logger, + ) from exc + return OpenAIAutoSwitchResponse( + enabled = enabled, + auto_unload_idle_seconds = idle_seconds, + idle_unload_active = get_auto_unload_idle_seconds() > 0, + ) + + +@router.get("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def get_openai_auto_switch_overrides( + current_subject: str = Depends(get_current_subject), +) -> ModelOverridesResponse: + return ModelOverridesResponse(overrides = get_model_overrides()) + + +@router.put("/openai-auto-switch/overrides", response_model = ModelOverridesResponse) +def update_openai_auto_switch_override( + payload: ModelOverridePayload, current_subject: str = Depends(get_current_subject) +) -> ModelOverridesResponse: + from core.inference.llama_server_args import validate_extra_args + try: + extra_args = validate_extra_args(payload.llama_extra_args) + set_model_override( + payload.model_id, + llama_extra_args = extra_args, + max_seq_length = payload.max_seq_length, + ) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid model launch override."), + event = "settings.update_model_override_failed", + log = logger, + ) from exc + return ModelOverridesResponse(overrides = get_model_overrides()) + + class PreviewLinkRotateResponse(BaseModel): rotated: bool = True diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4f131ad2f2..1da1c4f425 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -47,7 +47,7 @@ except ImportError: from utils.paths import resolve_dataset_path # Auth -from auth.authentication import get_current_subject +from auth.authentication import authenticated_via_api_key, get_current_subject from utils.utils import log_and_http_error @@ -114,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu @router.post("/start") async def start_training( - request: TrainingStartRequest, current_subject: str = Depends(get_current_subject) + request: TrainingStartRequest, + current_subject: str = Depends(get_current_subject), + via_api_key: bool = Depends(authenticated_via_api_key), ): """ Start a training job. @@ -125,6 +127,22 @@ async def start_training( try: logger.info(f"Starting training job with model: {request.model_name}") + # When Studio is driven as an inference API (API-key auth), refuse to start + # training while a request is in flight: training frees VRAM by unloading + # the chat model, which would kill the stream. The Studio UI (session auth) + # still starts training and coexists/frees VRAM as before. (A mixed UI+API + # session is not yet special-cased.) + if via_api_key is True: + from core.inference.llama_keepwarm import other_inference_request_count + if other_inference_request_count(current_request_counted = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Cannot start training over the API while an inference request is in " + "progress. Wait for it to finish, or start training from the Studio UI." + ), + ) + # No in-process ensure_transformers_version(): the subprocess # (worker.py) activates the correct version before importing ML libs. diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 23b90d7002..ba9f5b9cbc 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -1689,6 +1689,43 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]: conn.close() +def upsert_app_setting_map_entry( + key: str, entry_key: str, entry_value: dict[str, Any] | None +) -> dict[str, Any]: + """Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued + app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other + sub-entries cannot drop each other's updates.""" + conn = get_connection() + try: + conn.execute("BEGIN IMMEDIATE") + row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone() + current = _json_loads(row["value_json"], {}) if row else {} + if not isinstance(current, dict): + current = {} + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + now = datetime.now(timezone.utc).isoformat() + conn.execute( + """ + INSERT INTO app_settings (key, value_json, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(key) DO UPDATE SET + value_json = excluded.value_json, + updated_at = excluded.updated_at + """, + (key, json.dumps(current), now), + ) + conn.commit() + return current + except Exception: + conn.rollback() + raise + finally: + conn.close() + + def list_chat_settings() -> dict[str, Any]: conn = get_connection() try: diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py new file mode 100644 index 0000000000..7d2e2213b3 --- /dev/null +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -0,0 +1,3039 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in OpenAI /v1 model auto-switch: resolver, hook, and settings coercion. + +No GPU or llama-server: the backend and the load route are mocked, mirroring +tests/test_gguf_completion_usage.py. +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import LoadRequest +from core.inference import local_model_resolver as resolver +from utils import openai_auto_switch_settings as settings + + +class _FakeBackend: + def __init__( + self, + loaded_id = None, + hf_variant = None, + advertised_id = None, + ): + self.model_identifier = loaded_id + self.is_loaded = loaded_id is not None + self.hf_variant = hf_variant + self._openai_advertised_id = advertised_id + + +class _LoadRecorder: + """Stand-in for the load route: records calls and simulates a load.""" + + def __init__( + self, + backend, + fail = False, + ): + self.backend = backend + self.calls = [] + self.fail = fail + + async def __call__( + self, + request, + fastapi_request, + current_subject = None, + ): + self.calls.append(request) + if self.fail: + from fastapi import HTTPException + raise HTTPException(status_code = 503, detail = "load failed") + self.backend.model_identifier = request.model_path + self.backend.is_loaded = True + # Mirror _load_model_impl: a load advertises its own id until the + # auto-switch caller overwrites it with the repo id. + self.backend._openai_advertised_id = None + return None + + +def _wire(monkeypatch, *, enabled, resolves_to, backend, recorder): + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda _m: resolves_to) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + # Auto-switch loads via _load_model_impl (the /load route holds the lifecycle + # gate that auto-switch already owns, so it calls the impl directly). + monkeypatch.setattr(inference_route, "_load_model_impl", recorder) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + +def _run_hook(model = "some/model"): + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "tester")) + + +def test_flag_off_never_loads(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") + assert rec.calls == [] + + +def test_unknown_model_falls_through(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_already_loaded_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + # Case-insensitive match against the loaded identifier. + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/a-gguf", None, "unsloth/a-gguf"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/A-GGUF") + assert rec.calls == [] + + +def test_known_unloaded_model_switches_once(monkeypatch): + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert isinstance(req, LoadRequest) + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert backend.model_identifier == "unsloth/B-GGUF" + + +def test_concurrent_same_target_loads_once(monkeypatch): + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + + async def _race(): + await asyncio.gather( + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + inference_route._maybe_auto_switch_model("unsloth/B-GGUF", object(), "t"), + ) + + asyncio.run(_race()) + assert len(rec.calls) == 1 + + +def test_load_failure_propagates(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + rec = _LoadRecorder(backend, fail = True) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException): + _run_hook("unsloth/B-GGUF") + + +def test_same_repo_different_variant_switches(monkeypatch): + # Q4_K_M loaded, Q8_0 requested: a different quant must trigger a reload. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec.calls) == 1 + assert rec.calls[0].gguf_variant == "Q8_0" + + +def test_same_repo_same_variant_does_not_reload(monkeypatch): + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "q4_k_m", "unsloth/B-GGUF"), # case-insensitive + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF:Q4_K_M") + assert rec.calls == [] + + +def test_responses_endpoint_wires_auto_switch_before_dispatch(): + # The /v1/responses endpoint must invoke the auto-switch hook before either + # dispatcher so streaming requests switch too. Asserted on the source, which + # is immune to test-ordering effects on the shared inference module. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "_maybe_auto_switch_model" in src + hook_at = src.index("_maybe_auto_switch_model") + assert hook_at < src.index("_responses_stream") + assert hook_at < src.index("_responses_non_streaming") + + +def test_embeddings_endpoint_wires_auto_switch_before_loaded_check(): + # /v1/embeddings is model-bearing too, so it must auto-switch before the + # loaded-state gate. Asserted on the source for order-independence. + import inspect + + src = inspect.getsource(inference_route.openai_embeddings) + assert "_auto_switch_from_request_body" in src + assert src.index("_auto_switch_from_request_body") < src.index("is_loaded") + + +def test_count_tokens_endpoint_wires_auto_switch_before_loaded_check(): + # The Anthropic token-count endpoint must count with the requested model. + import inspect + + src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "_maybe_auto_switch_model" in src + assert src.index("_maybe_auto_switch_model") < src.index("is_loaded") + + +def test_openai_compat_routes_bound_to_handlers_with_auth(): + # Inserting a helper between a @router.post decorator and its handler silently + # rebinds the route to the helper and drops its auth dependency (this happened to + # /messages/count_tokens). The source-inspection tests above miss it because they + # call the handler directly. Lock the path -> (handler, auth) mapping at the route + # level so any decorator/handler split is caught. + expected = { + ("POST", "/chat/completions"): "openai_chat_completions", + ("POST", "/completions"): "openai_completions", + ("POST", "/embeddings"): "openai_embeddings", + ("POST", "/responses"): "openai_responses", + ("POST", "/messages"): "anthropic_messages", + ("POST", "/messages/count_tokens"): "anthropic_count_tokens", + ("POST", "/audio/generate"): "generate_audio", + ("GET", "/models"): "openai_list_models", + ("GET", "/models/{model_id:path}"): "openai_retrieve_model", + } + seen = {} + for r in inference_route.router.routes: + path = getattr(r, "path", None) + endpoint = getattr(r, "endpoint", None) + if path is None or endpoint is None: + continue + for method in getattr(r, "methods", None) or (): + seen[(method, path)] = r + for key, handler in expected.items(): + assert key in seen, f"route {key} is not registered" + route = seen[key] + assert ( + route.endpoint.__name__ == handler + ), f"{key} bound to {route.endpoint.__name__}, expected {handler}" + deps = [d.call.__name__ for d in route.dependant.dependencies] + assert "get_current_subject" in deps, f"{key} lost its auth dependency" + + +# ── resolver ──────────────────────────────────────────────────────── + + +def test_local_gguf_entry_filters_non_gguf_and_recurses(tmp_path): + from types import SimpleNamespace + + # Transformers/safetensors folder: not a GGUF, must be rejected. + tf = tmp_path / "tf-model" + tf.mkdir() + (tf / "config.json").write_text("{}") + (tf / "model.safetensors").write_text("x") + assert resolver._local_gguf_entry("tf", SimpleNamespace(path = str(tf))) is None + + # Standalone .gguf file: an entry with no quant sub-selection. + bare = tmp_path / "x.gguf" + bare.write_text("x") + e = resolver._local_gguf_entry("x", SimpleNamespace(path = str(bare))) + assert e is not None and e.variants == () + + # HF-cache snapshots with a quant subdir (the nested layout the previous + # shallow glob missed): must still be detected. + repo = tmp_path / "models--org--repo" + (repo / "snapshots" / "abc" / "BF16").mkdir(parents = True) + (repo / "snapshots" / "abc" / "BF16" / "model-BF16.gguf").write_text("x") + e2 = resolver._local_gguf_entry("org/repo", SimpleNamespace(path = str(repo))) + assert e2 is not None and e2.variants + + +def test_local_gguf_entry_rejects_standalone_mmproj(tmp_path): + # Codex P2: _scan_models_dir's standalone-.gguf pass emits an entry for a + # bare mmproj projector (it only filters mmproj inside directory scans). A + # projector is not a servable model, so the resolver must reject it or + # /v1/models advertises it and a switch could load it over the real weights. + from types import SimpleNamespace + + proj = tmp_path / "mmproj-F16.gguf" + proj.write_text("x") + assert resolver._local_gguf_entry("p", SimpleNamespace(path = str(proj))) is None + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(proj), path = str(proj))) is False + + +def _entry(loader_id, *variants): + # load_path == loader_id for tests; production stores a concrete local path. + return resolver._LocalGgufEntry(loader_id, loader_id, tuple(variants)) + + +def test_resolver_matches_and_splits_variant(monkeypatch): + monkeypatch.setattr( + resolver, + "_build_index", + lambda: {"unsloth/b-gguf": _entry("unsloth/B-GGUF", "UD-Q5_K_XL", "Q4_K_M")}, + ) + resolver._scan = (0.0, {}) # force a rescan + # A requested variant present on disk resolves (case-insensitive). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:ud-q5_k_xl") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A bare id resolves to a concrete local quant, never a remote one. + assert resolver.resolve_local_gguf("unsloth/B-GGUF") == ( + "unsloth/B-GGUF", + "UD-Q5_K_XL", + "unsloth/B-GGUF", + ) + # A variant that is not on disk must not resolve (no remote download). + assert resolver.resolve_local_gguf("unsloth/B-GGUF:Q8_0") is None + assert resolver.resolve_local_gguf("totally/unknown") is None + assert resolver.resolve_local_gguf("") is None + + +def test_resolver_failsafe_on_internal_error(monkeypatch): + # Resolution is best-effort: any internal failure must fall through to None + # so the request still serves the loaded model instead of 500-ing. The hook + # calls resolve_local_gguf without its own guard, so the guard lives here. + def boom(): + raise RuntimeError("scan blew up") + + monkeypatch.setattr(resolver, "_build_index", boom) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("unsloth/B-GGUF") is None + + +def test_resolver_nonstring_model_is_failsafe(): + # /v1/completions and /v1/embeddings pass body.get("model") straight through, + # so a non-string must not raise on .strip(). + assert resolver.resolve_local_gguf(123) is None + assert resolver.resolve_local_gguf({"a": 1}) is None + assert resolver.resolve_local_gguf(None) is None + + +def test_resolver_exact_id_with_colon_wins(monkeypatch): + # A local id that itself contains a colon (e.g. a Windows path) must match + # exactly rather than being split at the drive-letter colon. + win = r"C:\models\foo.gguf" + monkeypatch.setattr(resolver, "_build_index", lambda: {win.lower(): _entry(win)}) + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(win) == (win, None, win) + + +# ── settings coercion ─────────────────────────────────────────────── + + +def test_setting_coercion(): + assert settings._coerce_bool("on") is True + assert settings._coerce_bool("off") is False + assert settings._coerce_bool("garbage") is None + assert settings._coerce_int("5") == 5 + assert settings._coerce_int(-3) == 0 + assert settings._coerce_int("nope") is None + + +# ── idle keep-warm ────────────────────────────────────────────────── + + +def test_idle_loop_does_not_unload_freshly_loaded_model(monkeypatch): + # Server idle far longer than the TTL, then a model is loaded: the load + # transition stamps activity so the next poll must not unload it. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 1) + kw._inflight = 0 + kw._last_active = time.monotonic() - 3600 + + unloads = [] + backend = _FakeBackend("unsloth/Fresh-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch): + # The headline behavior (the other idle tests only cover the negative paths): + # with nothing in flight and the TTL elapsed, the loop frees the GGUF exactly + # once and records its identity so a later alias request can reload that variant. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005) + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + kw._last_unloaded_model = None + + unloads = [] + backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M") + + def _unload(): + unloads.append(1) + backend.is_loaded = False # a real unload clears the slot + + backend.unload_model = _unload + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.02)) + await asyncio.sleep(0.2) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [1] # freed once, not repeatedly + stash = kw.get_last_unloaded_model() + assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M" + + +def test_audio_generate_is_tracked_as_inference_path(): + # Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so + # the keep-warm middleware must count it as in-flight inference. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/chat/completions") is True + assert _is_inference_path("/api/inference/models/list") is False + + +def test_idle_loop_does_not_unload_while_request_inflight(monkeypatch): + # An in-flight request (inflight > 0) must protect the model from unload + # even when it has been idle by wall-clock past the TTL. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.01) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_last_active", time.monotonic() - 3600) + + unloads = [] + backend = _FakeBackend("unsloth/Active-GGUF") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.08) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + assert unloads == [] + + +# ── per-model launch overrides ────────────────────────────────────── + + +def test_auto_switch_applies_model_override(monkeypatch): + # A configured model loads with its saved launch flags, not bare defaults. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, + "get_model_override", + lambda model_id: {"llama_extra_args": ["--n-gpu-layers", "20"], "max_seq_length": 4096}, + ) + + _run_hook("unsloth/B-GGUF") + assert len(rec.calls) == 1 + req = rec.calls[0] + assert req.model_path == "unsloth/B-GGUF" + assert req.gguf_variant == "Q4_K_M" + assert req.llama_extra_args == ["--n-gpu-layers", "20"] + assert req.max_seq_length == 4096 + + +def test_auto_switch_applies_partial_override(monkeypatch): + # Only llama_extra_args is configured: it is applied, max_seq_length stays default. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q4_K_M", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr( + settings, "get_model_override", lambda model_id: {"llama_extra_args": ["--flash-attn"]} + ) + + _run_hook("unsloth/B-GGUF") + req = rec.calls[0] + assert req.llama_extra_args == ["--flash-attn"] + assert req.max_seq_length == 0 # untouched default + + +def _mock_override_store(monkeypatch): + """Back the override read + atomic-merge write with an in-memory dict.""" + import storage.studio_db as db + + store = {} + + def _merge_entry(key, entry_key, entry_value): + current = dict(store.get(key) or {}) + if entry_value: + current[entry_key] = entry_value + else: + current.pop(entry_key, None) + store[key] = current + return current + + monkeypatch.setattr(db, "upsert_app_setting_map_entry", _merge_entry) + monkeypatch.setattr(db, "get_app_setting", lambda k, default = None: store.get(k, default)) + settings._cache.clear() + return store + + +def test_model_override_roundtrip(monkeypatch): + _mock_override_store(monkeypatch) + + settings.set_model_override( + "unsloth/B-GGUF", llama_extra_args = ["--n-gpu-layers", "20"], max_seq_length = 4096 + ) + assert settings.get_model_override("unsloth/B-GGUF") == { + "llama_extra_args": ["--n-gpu-layers", "20"], + "max_seq_length": 4096, + } + # An override with no fields removes the entry rather than storing an empty one. + settings.set_model_override("unsloth/B-GGUF", llama_extra_args = [], max_seq_length = None) + assert settings.get_model_override("unsloth/B-GGUF") == {} + assert settings.get_model_overrides() == {} + + +def test_override_route_rejects_managed_flag_and_removes(monkeypatch): + import routes.settings as settings_route + from fastapi import HTTPException + + _mock_override_store(monkeypatch) + + # A managed/denylisted llama-server flag is rejected with 400, not 500. + bad = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--port", "1234"] + ) + with pytest.raises(HTTPException) as excinfo: + settings_route.update_openai_auto_switch_override(bad, "tester") + assert excinfo.value.status_code == 400 + + # A valid override is stored, then an empty payload removes it through the route. + ok = settings_route.ModelOverridePayload( + model_id = "unsloth/B-GGUF", llama_extra_args = ["--flash-attn"], max_seq_length = 4096 + ) + resp = settings_route.update_openai_auto_switch_override(ok, "tester") + assert resp.overrides["unsloth/B-GGUF"]["max_seq_length"] == 4096 + assert "llama_extra_args" in resp.overrides["unsloth/B-GGUF"] + + empty = settings_route.ModelOverridePayload(model_id = "unsloth/B-GGUF") + resp2 = settings_route.update_openai_auto_switch_override(empty, "tester") + assert "unsloth/B-GGUF" not in resp2.overrides + + +def test_model_override_rejects_zero_max_seq_length(): + # 0 is not a valid sequence length and the setter drops a falsy value, so the + # payload must reject it at the boundary instead of accepting then discarding it. + import pydantic + import routes.settings as settings_route + + with pytest.raises(pydantic.ValidationError): + settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 0) + assert settings_route.ModelOverridePayload(model_id = "x", max_seq_length = 1).max_seq_length == 1 + + +def test_update_openai_auto_switch_writes_both_keys_in_one_transaction(monkeypatch): + # The PUT must persist enabled + idle in a single upsert so a settings write can't + # leave one key updated and the other stale. + import routes.settings as settings_route + import storage.studio_db as db + from utils.openai_auto_switch_settings import ( + AUTO_UNLOAD_IDLE_SETTING_KEY, + OPENAI_AUTO_SWITCH_SETTING_KEY, + ) + + calls = [] + + def _capture(mapping): + calls.append(dict(mapping)) + return {} + + monkeypatch.setattr(db, "upsert_app_settings", _capture) + settings._cache.clear() + + payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_idle_seconds = 120) + resp = settings_route.update_openai_auto_switch(payload, "tester") + assert resp.enabled is True and resp.auto_unload_idle_seconds == 120 + assert len(calls) == 1 # one transaction, not two + written = calls[0] + assert written.get(OPENAI_AUTO_SWITCH_SETTING_KEY) is True + assert written.get(AUTO_UNLOAD_IDLE_SETTING_KEY) == 120 + + +def test_settings_report_idle_unload_active_when_env_backed(monkeypatch): + # Codex P2: with UNSLOTH_MODEL_IDLE_TTL driving idle-unload while the toggle is + # off, the settings response must report idle_unload_active so the UI shows the + # feature as active via env rather than "needs enable". + import routes.settings as settings_route + + monkeypatch.setattr(settings_route, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(settings_route, "get_stored_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr( + settings_route, "get_auto_unload_idle_seconds", lambda: 600 + ) # effective > 0 + resp = settings_route.get_openai_auto_switch("tester") + assert resp.enabled is False and resp.idle_unload_active is True + # Effective TTL 0 (off, nothing env-backed) -> not active. + monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0) + assert settings_route.get_openai_auto_switch("tester").idle_unload_active is False + + +# ── /v1/models discovery ──────────────────────────────────────────── + + +def test_v1_models_retrieve_is_case_insensitive(monkeypatch): + # The resolver lowercases its index, so a retrieve that differs only in case + # from a catalog id must still hit (200), not 404. Guards the .lower() compare + # in openai_retrieve_model against a silent revert. (The full local catalog is + # main's #6519; only the loaded fast-path is exact, the catalog loop is lenient.) + from fastapi import HTTPException + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + + async def _catalog(): + return [ + {"id": "unsloth/A-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + {"id": "unsloth/B-GGUF", "object": "model", "created": 1, "owned_by": "local"}, + ] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _catalog) + + # A catalog id retrieved with different casing still resolves. + obj = asyncio.run(inference_route.openai_retrieve_model("unsloth/a-gguf", "tester")) + assert obj["id"] == "unsloth/A-GGUF" + # A truly unknown id still 404s. + with pytest.raises(HTTPException) as unknown: + asyncio.run(inference_route.openai_retrieve_model("totally/unknown", "tester")) + assert unknown.value.status_code == 404 + + +# ── hardening: hidden models, idle/enabled coupling, count_tokens keep-warm ── + + +def test_index_excludes_hidden_models(tmp_path, monkeypatch): + # The llama.cpp validation probe and RAG embedding weights are hidden from + # Studio's pickers; they must never become auto-switch targets. + from types import SimpleNamespace + import routes.models as models_route + + normal = tmp_path / "normal-Q4_K_M.gguf" + normal.write_bytes(b"x" * 32) + probe = tmp_path / "stories260K.gguf" # llama.cpp install-validation probe + probe.write_bytes(b"x" * 32) + + def _info(mid, path): + return SimpleNamespace(id = mid, path = str(path), model_id = mid, display_name = mid) + + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda *a, **k: [_info("org/Normal-GGUF", normal), _info("ggml-org/models", probe)], + ) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + resolver._scan = (0.0, {}) + + index = resolver._index() + assert "org/normal-gguf" in index # keys are normalized to lowercase + assert "ggml-org/models" not in index + # And the hidden probe cannot be auto-switched to by name. + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf("ggml-org/models") is None + + +def test_idle_disabled_when_auto_switch_off(monkeypatch): + # "Off means unchanged": a stored idle TTL must report 0 while auto-switch is + # off, so the idle loop and keep-warm middleware can never unload the model. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 60} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 60 + + +def test_count_tokens_is_tracked_as_inference_path(): + # count_tokens counts via the loaded tokenizer, so idle-unload must not pull + # the model out from under it; it has to be a tracked in-flight path. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/v1/messages/count_tokens") is True + assert _is_inference_path("/api/inference/messages/count_tokens") is True + assert _is_inference_path("/v1/messages") is True + + +# ── review follow-ups: bare-id reuse, responses order, in-flight tracking ── + + +def test_bare_id_tolerates_any_loaded_variant(monkeypatch): + # Repo already loaded as Q4_K_M; a BARE request for the same repo (resolver + # picks the largest local quant, Q8_0) must NOT reload a different quant. + backend = _FakeBackend("unsloth/B-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("unsloth/B-GGUF") # bare, no :VARIANT + assert rec.calls == [] + # An explicit :VARIANT request still honors the quant (reloads to Q8_0). + rec2 = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("unsloth/B-GGUF", "Q8_0", "unsloth/B-GGUF"), + backend = backend, + recorder = rec2, + ) + _run_hook("unsloth/B-GGUF:Q8_0") + assert len(rec2.calls) == 1 + + +def test_responses_hook_runs_after_input_validation(): + # A request that 400s on empty input must not have triggered a model load, + # so the auto-switch hook must come after the input-validation guard. + import inspect + + src = inspect.getsource(inference_route.openai_responses) + assert "No input provided" in src + assert src.index("No input provided") < src.index("_maybe_auto_switch_model") + + +def test_responses_system_only_rejected_before_switch(monkeypatch): + # Codex P2: instructions-only input normalises to a lone system message, which + # passes the empty-input check; it must 400 before the switch so an invalid + # Responses request can't evict the resident model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch a system-only Responses request") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", instructions = "be helpful", input = "") + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + + +def test_keepwarm_tracks_inflight_when_enabled_even_if_idle_zero(monkeypatch): + # In-flight must be counted whenever auto-switch is on, even with idle TTL 0, + # so enabling idle mid-stream cannot unload an in-flight request. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + kw._inflight = 0 + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # counted despite idle TTL being 0 + assert kw._inflight == 0 # balanced after completion + + +# ── review follow-ups: OFF-state body, swap guard, alias reload, always-track ── + + +def _bad_body_request(): + import json as _json + class _BadReq: + async def json(self): + raise _json.JSONDecodeError("expecting value", "", 0) + + return _BadReq() + + +def test_completions_malformed_body_503_not_500_when_unloaded(monkeypatch): + # OFF + nothing loaded + unparseable body must still 503 (pre-feature + # behavior), not 500 from the early body read. + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_embeddings_malformed_body_503_not_500_when_unloaded(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend(None) + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_bad_body_request(), "tester")) + assert exc.value.status_code == 503 + + +def test_non_string_model_falls_through_without_error(monkeypatch): + # A non-string model (e.g. {"model": 123} on a raw-body endpoint) must be + # treated as absent, never raising in the membership checks, even when a stash + # exists from idle-unload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + asyncio.run(inference_route._maybe_auto_switch_model(123, object(), "tester")) + assert rec.calls == [] # no load, no TypeError + + +def test_anthropic_validates_max_tokens_before_auto_switch(): + # An Anthropic request missing max_tokens must 400 before the hook runs, so an + # invalid request never triggers a model load. Asserted on the source order. + import inspect + + src = inspect.getsource(inference_route.anthropic_messages) + assert "_maybe_auto_switch_model" in src + assert src.index("max_tokens: field required") < src.index("_maybe_auto_switch_model") + + +def test_alias_reloads_model_freed_by_idle_unload_with_quant(monkeypatch): + # After idle-unload frees the model, an unknown/alias name (resolves to None) + # reloads what was freed, including the exact quant, instead of 503-ing. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the backend + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", "Q4_K_M")) + _run_hook("gpt-4o-mini") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "unsloth/A-GGUF" + assert rec.calls[0].gguf_variant == "Q4_K_M" # exact freed quant restored + + +def test_alias_does_not_reload_when_model_already_loaded(monkeypatch): + # The reload only triggers on an empty backend; with something loaded, an + # unknown name still falls through (drop-in) without resurrecting the stash. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("unsloth/B-GGUF") + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_last_unloaded_model", ("unsloth/A-GGUF", None)) + _run_hook("gpt-4o-mini") + assert rec.calls == [] + + +def test_idle_loop_does_not_unload_while_request_pending(monkeypatch): + # A request that has marked itself pending (waiting on the unload gate) but not + # yet started must keep the idle loop from unloading the model. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 0.0) # far past any TTL + kw._note_pending() + try: + assert kw._is_idle(1.0) is False # pending request blocks unload + finally: + kw._note_unpending() + assert kw._is_idle(1.0) is True # cleared once it is no longer pending + + +def test_keepwarm_tracks_inflight_even_when_auto_switch_off(monkeypatch): + # A stream that starts while the feature is OFF must still be counted, so + # enabling idle-unload mid-stream cannot unload it. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"ok", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 1 # tracked despite the feature being off + assert kw._inflight == 0 + + +def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch, tmp_path): + # _build_index must scan the same roots the model picker lists, else a model + # the UI shows is silently served as the loaded one. Verify each is consulted. + from pathlib import Path + import routes.models as models_route + from utils import paths as upaths + import storage.studio_db as studio_db + + scanned = [] + monkeypatch.setattr( + models_route, + "_scan_models_dir", + lambda d, limit = None: scanned.append(("models", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_hf_cache", + lambda d: scanned.append(("hf", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr( + models_route, + "_scan_lmstudio_dir", + lambda d: scanned.append(("lm", str(Path(d).resolve()))) or [], + ) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path / "active") + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(upaths, "legacy_hf_cache_dir", lambda: tmp_path / "legacy") + monkeypatch.setattr(upaths, "hf_default_cache_dir", lambda: tmp_path / "default") + monkeypatch.setattr(upaths, "lmstudio_model_dirs", lambda: [tmp_path / "lmstudio"]) + monkeypatch.setattr( + studio_db, "list_scan_folders", lambda: [{"path": str(tmp_path / "custom")}] + ) + for sub in ("active", "legacy", "default", "lmstudio", "custom"): + (tmp_path / sub).mkdir() + + resolver._build_index() + + hf = {p for k, p in scanned if k == "hf"} + lm = {p for k, p in scanned if k == "lm"} + assert str((tmp_path / "legacy").resolve()) in hf + assert str((tmp_path / "default").resolve()) in hf + assert str((tmp_path / "custom").resolve()) in hf + assert str((tmp_path / "lmstudio").resolve()) in lm + + +# ── gemini round: list-body 400, non-POST not tracked ── + + +def _json_body_request(payload): + class _Req: + async def json(self): + return payload + + return _Req() + + +def test_completions_list_body_is_400_not_500(monkeypatch): + # A valid JSON non-dict body (e.g. a list) on a loaded backend is a clean 400, + # not a 500 from body.get(...). + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") # loaded + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_completions(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_embeddings_list_body_is_400_not_500(monkeypatch): + from fastapi import HTTPException + + backend = _FakeBackend("unsloth/A-GGUF") + _wire( + monkeypatch, + enabled = False, + resolves_to = None, + backend = backend, + recorder = _LoadRecorder(backend), + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request([]), "tester")) + assert exc.value.status_code == 400 + + +def test_middleware_ignores_non_post(monkeypatch): + # CORS preflight (OPTIONS) on an inference path must not be tracked as in-flight. + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + seen = {} + + async def app(scope, receive, send): + seen["inflight"] = kw._inflight + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b"", "more_body": False}) + + async def drive(): + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(_m): + pass + + scope = {"type": "http", "path": "/v1/chat/completions", "method": "OPTIONS", "headers": []} + await kw.LlamaKeepWarmMiddleware(app)(scope, receive, send) + + asyncio.run(drive()) + assert seen["inflight"] == 0 # OPTIONS not counted + assert kw._inflight == 0 + + +# ── review round 4: swap guard, idle variant identity, load-by-path, stash clear ── + + +def test_auto_switch_refuses_when_another_inference_is_active(monkeypatch): + # A cross-model swap must 409 (not kill) while another inference request is in + # flight; the requesting call itself is excluded from the count. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # this request + another active one + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_auto_switch_swaps_when_only_caller_is_active(monkeypatch): + # Only the caller is in flight: nothing else to protect, so the swap proceeds. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", None, "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/B-GGUF") + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/p/B" # concrete local path, not the repo id + + +def test_idle_loop_resets_timer_for_same_repo_different_variant(monkeypatch): + # Same repo, different quant counts as a fresh model: the idle timer resets, so + # the new variant is not unloaded before one TTL of its own. + import time + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.05) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + + unloads = [] + backend = _FakeBackend("org/model-GGUF", hf_variant = "Q4_K_M") + backend.unload_model = lambda: unloads.append(1) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + + async def _drive(): + task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01)) + await asyncio.sleep(0.03) + assert unloads == [] + kw._last_active = time.monotonic() - 60 # force idle + backend.hf_variant = "Q8_0" # same id, new quant -> fresh identity + await asyncio.sleep(0.03) + assert unloads == [] # timer reset by the variant change, not unloaded + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(_drive()) + + +def test_generate_stream_is_tracked_as_inference_path(): + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/api/inference/generate/stream") is True + assert _is_inference_path("/api/inference/audio/generate") is True + assert _is_inference_path("/v1/responses") is True + + +def test_successful_manual_load_clears_last_unloaded_stash(): + from core.inference import llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_loaded() + assert kw.get_last_unloaded_model() is None + + +def test_hf_cache_entry_loads_from_local_snapshot_path(tmp_path): + # An HF-cache repo resolves to its on-disk snapshot dir, so /load takes the + # local branch (no repo-id download). loader_id stays the repo id. + from types import SimpleNamespace + + repo = tmp_path / "models--org--Repo" + snap = repo / "snapshots" / "abc123" + snap.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"GGUF stub") + + entry = resolver._local_gguf_entry("org/Repo", SimpleNamespace(id = "org/Repo", path = str(repo))) + assert entry is not None + assert entry.loader_id == "org/Repo" # advertised id unchanged + assert "snapshots" in entry.load_path # loads from the concrete snapshot dir + assert entry.load_path != "org/Repo" # never the bare repo id + assert entry.variants # quant detected on disk + + +# ── review round 5: concurrent-swap, repo-id identity, /v1/models id, gate, 503 ── + + +def test_already_loaded_by_repo_id_is_not_reswapped(monkeypatch): + # A model loaded normally has model_identifier == repo id, but the resolver + # returns the concrete load path. A request for that repo must count as already + # serving (no reload, no 409) even with another inference active. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/Repo-GGUF", hf_variant = "Q4_K_M") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/cache/models--org--Repo-GGUF/snapshots/abc", "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + _run_hook("org/Repo-GGUF:Q4_K_M") # exact quant + _run_hook("org/Repo-GGUF") # bare id + assert rec.calls == [] + + +def test_auto_switch_advertises_repo_id_after_load(monkeypatch): + # After a load-by-path, the backend advertises the repo id (override key), not + # the concrete path, so /v1/models and the idle stash stay name-based. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B-snapshot", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + _run_hook("org/B-GGUF:Q8_0") + assert rec.calls[0].model_path == "/p/B-snapshot" # loaded by concrete path + assert backend._openai_advertised_id == "org/B-GGUF" # advertised by repo id + + +def test_already_serving_by_path_records_advertised_alias(monkeypatch): + # Codex P2: a model loaded by local path and requested via an advertised alias + # that resolves to the same path is already serving (no reload), but /v1/models + # and responses would report the path basename and list the alias as loaded:false + # unless the alias is recorded as the advertised id on the already-serving return. + path = "/cache/models--org--Repo-GGUF/snapshots/abc" + backend = _FakeBackend(path, hf_variant = "Q4_K_M") # loaded by path, no advertised id + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = (path, "Q4_K_M", "org/Repo-GGUF"), + backend = backend, + recorder = rec, + ) + assert backend._openai_advertised_id is None + _run_hook("org/Repo-GGUF:Q4_K_M") + assert rec.calls == [] # already serving -> no reload + assert backend._openai_advertised_id == "org/Repo-GGUF" # alias now recorded + + +def test_streaming_responses_uses_advertised_id_helper(): + # Codex P2: streamed /v1/responses envelopes must derive the model id from + # _llama_public_model_id (which prefers _openai_advertised_id), not the raw + # model_identifier. After an auto-switch to a cached HF GGUF the identifier is + # the snapshot path while the repo id lives in _openai_advertised_id, so the raw + # form would stream a snapshot basename while /v1/models, chat, and non-streaming + # responses report the repo id. + import inspect + + src = inspect.getsource(inference_route._responses_stream) + assert "_clean_model = _llama_public_model_id(llama_backend" in src + assert 'public_model_id(getattr(llama_backend, "model_identifier"' not in src + + +def test_concurrent_same_target_requests_load_once(monkeypatch): + # Two concurrent requests for the same unloaded model must load once, not each + # 409 the other. Simulate the second request already waiting (registered) while + # the first runs the hook with _inflight counting both. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # both same-target requests counted + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/B-GGUF", "Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_swap_still_refused_when_other_request_targets_different_model(monkeypatch): + # A concurrent request heading to a different target still blocks the swap: the + # same-target exclusion must not swallow a genuinely conflicting request. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) + monkeypatch.setattr(kw, "_pending", 0) + inference_route._note_switch_waiter(inference_route._switch_key("org/C-GGUF", "Q4_K_M"), 1) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] + + +def test_v1_models_advertises_repo_id_not_load_path(monkeypatch): + # /v1/models must report the advertised repo id, never the host load path. + from types import SimpleNamespace + + llama = _FakeBackend("/cache/models--org--Repo/snapshots/abc") + llama._openai_advertised_id = "org/Repo-GGUF" + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr( + inference_route, "get_inference_backend", lambda: SimpleNamespace(active_model_name = None) + ) + objects = inference_route._openai_model_objects() + assert [o["id"] for o in objects] == ["org/Repo-GGUF"] + + +def test_idle_alias_reload_preserves_override_via_advertised_id(monkeypatch): + # The idle stash carries (load_path, quant, advertised_id). An alias reload must + # look up the override by the advertised repo id, not the concrete load path, + # so the user's saved launch flags survive the unload/reload. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + overrides = {"org/A-GGUF": {"max_seq_length": 8192}} + monkeypatch.setattr(settings, "get_model_override", lambda mid: overrides.get(mid, {})) + _run_hook("gpt-4o-mini") + assert rec.calls[0].model_path == "/cache/snap/A" # reloads the freed path + assert rec.calls[0].gguf_variant == "Q4_K_M" + assert rec.calls[0].max_seq_length == 8192 # override keyed by repo id, not path + + +def test_load_route_holds_lifecycle_gate(monkeypatch): + # Lock the manual /load gate against silent revert: the route must wrap the + # load in inference_lifecycle_gate so idle-unload can't fire mid-load. + import inspect + + src = inspect.getsource(inference_route.load_model) + assert "inference_lifecycle_gate" in src + assert "_load_model_impl" in src + + +def _anthropic_payload(max_tokens = None): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "claude-x", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + ) + + +def test_anthropic_503_when_unloaded_and_auto_switch_off(monkeypatch): + # Default-off parity: unloaded backend + auto-switch off 503s before the + # max_tokens 400, exactly as the pre-feature endpoint did. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 503 + + +def test_anthropic_400_when_auto_switch_on_and_max_tokens_missing(monkeypatch): + # With auto-switch on, request-shape validation runs first: a missing + # max_tokens still 400s before any load is attempted. + from fastapi import HTTPException + + backend = _FakeBackend(None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(_anthropic_payload(), object(), "tester")) + assert exc.value.status_code == 400 + + +# ── review round 6: concurrency ordering, external untrack, unload gate, ids ── + + +def test_pending_same_target_request_does_not_force_409(monkeypatch): + # A second same-target request blocked in the middleware (pending, not yet + # generating) must not make the first request 409: pending is excluded. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 1) # just the caller + monkeypatch.setattr(kw, "_pending", 1) # second request blocked in middleware + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_concurrent_same_target_loads_once_while_other_still_resolving(monkeypatch): + # The real middleware counts a concurrent same-model request as in-flight + # before it resolves and registers a target waiter. The raw-request waiter, + # registered before resolve, must still exclude it so the first request loads. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # caller + a still-resolving twin + monkeypatch.setattr(kw, "_pending", 0) + # The twin has only registered its raw requested model (not yet a target waiter). + inference_route._note_request_waiter(inference_route._request_waiter_key("org/B-GGUF:Q8_0"), 1) + _run_hook("org/B-GGUF:Q8_0") + assert len(rec.calls) == 1 # loads once, no 409 + + +def test_external_untrack_decrements_inflight_and_is_idempotent(): + from core.inference import llama_keepwarm as kw + + kw._inflight = 2 + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 1 + assert scope.get(kw._UNTRACKED_SCOPE_KEY) is True + kw.untrack_current_request(scope) # idempotent: no further decrement + assert kw._inflight == 1 + kw._inflight = 0 + + +def test_manual_unload_interrupts_even_while_inference_active(monkeypatch): + # A manual /unload is a deliberate action: it tears down immediately even with + # a request in flight (only the automatic idle loop defers). No 409. + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + backend = _FakeBackend("org/A-GGUF") + backend.is_active = True + backend.unload_model = lambda: setattr(backend, "is_loaded", False) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + monkeypatch.setattr(kw, "_inflight", 1) # another request streaming + monkeypatch.setattr(kw, "_pending", 0) + resp = asyncio.run( + inference_route.unload_model(UnloadRequest(model_path = "org/A-GGUF"), "tester") + ) + assert resp.status == "unloaded" + assert not backend.is_loaded # torn down despite the active request + + +def test_auto_switch_refuses_when_unsloth_stream_active(monkeypatch): + # The GGUF slot is empty but an Unsloth model is streaming (counted in-flight). + # _load_model_impl would unload it, so auto-switch must 409, not only when a + # GGUF is loaded. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # no GGUF loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(kw, "_inflight", 2) # an Unsloth stream + this request + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + _run_hook("org/B-GGUF:Q8_0") + assert exc.value.status_code == 409 + assert rec.calls == [] # the active Unsloth model is not torn down + + +def test_public_model_id_prefers_advertised_over_path(): + backend = _FakeBackend("/cache/models--org--Repo/snapshots/abc/model.gguf") + backend._openai_advertised_id = "org/Repo-GGUF" + # The advertised repo id from an auto-switch load wins. + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend._openai_advertised_id = None + # No advertised id: the identifier is cleaned to a public id (delegates to + # public_model_id), never the raw on-disk .gguf path. + cleaned = inference_route._llama_public_model_id(backend) + assert cleaned and "/cache/" not in cleaned and not cleaned.endswith(".gguf") + # An already-clean repo id passes through unchanged. + backend.model_identifier = "org/Repo-GGUF" + assert inference_route._llama_public_model_id(backend) == "org/Repo-GGUF" + backend.model_identifier = None + assert inference_route._llama_public_model_id(backend, "req") == "req" + + +def test_chat_validates_non_system_message_before_auto_switch(): + # A system-only chat must be rejected before the hook so an invalid request + # never swaps the resident model. Asserted on source order. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("At least one non-system message is required.") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_chat_untracks_external_provider_before_proxy(): + # The external-provider branch must untrack the request before proxying so its + # stream can't block a concurrent local auto-switch. + import inspect + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("untrack_current_request") < src.index("_proxy_to_external_provider") + + +# ── round 7: API-initiated training defers to active inference, UI does not ── + + +def test_authenticated_via_api_key_detects_key_vs_session(): + from fastapi.security import HTTPAuthorizationCredentials + from auth.authentication import authenticated_via_api_key, API_KEY_PREFIX + + key = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = API_KEY_PREFIX + "abc") + jwt = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = "eyJhbGciOiJ.session") + assert asyncio.run(authenticated_via_api_key(key)) is True + assert asyncio.run(authenticated_via_api_key(jwt)) is False + + +def _training_request(): + from models.training import TrainingStartRequest + return TrainingStartRequest( + model_name = "unsloth/test", training_type = "LoRA/QLoRA", format_type = "alpaca" + ) + + +def test_api_training_refused_while_inference_active(monkeypatch): + # API-key caller: training is refused with 409 while a request streams, so it + # can't free VRAM by unloading the chat model out from under the stream. + from fastapi import HTTPException + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + with pytest.raises(HTTPException) as exc: + asyncio.run( + training_route.start_training( + _training_request(), current_subject = "t", via_api_key = True + ) + ) + assert exc.value.status_code == 409 + + +def test_ui_training_not_blocked_by_active_inference(monkeypatch): + # UI (session auth) caller: the API guard is skipped, so training proceeds past + # it even with inference active (here it hits the normal already-active path). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + import routes.training as training_route + + monkeypatch.setattr(kw, "_inflight", 1) + monkeypatch.setattr(kw, "_pending", 0) + fake = SimpleNamespace(is_training_active = lambda: True, current_job_id = "job-1") + monkeypatch.setattr(training_route, "get_training_backend", lambda: fake) + resp = asyncio.run( + training_route.start_training(_training_request(), current_subject = "t", via_api_key = False) + ) + assert resp.status == "error" and "already" in (resp.error or "").lower() + + +# ── UNSLOTH_MODEL_IDLE_TTL env override (borrowed from PR 6517) ── + + +def test_env_idle_ttl_standalone_when_no_stored_value(monkeypatch): + # With nothing stored, the env var enables idle-unload even while auto-switch + # is off (headless/ops default), and the UI reader reflects it. + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) # nothing stored + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 600 + assert settings.get_stored_auto_unload_idle_seconds() == 600 + + +def test_stored_idle_value_overrides_env_and_stays_gated(monkeypatch): + # An explicit stored value wins over the env default and remains gated on the + # auto-switch toggle. + store = {settings.AUTO_UNLOAD_IDLE_SETTING_KEY: 30} + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d)) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "600") + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert settings.get_auto_unload_idle_seconds() == 30 # stored wins, not env + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + assert settings.get_auto_unload_idle_seconds() == 0 # explicit value still gated off + + +def test_env_idle_ttl_invalid_is_ignored(monkeypatch): + monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: d) + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + monkeypatch.setenv("UNSLOTH_MODEL_IDLE_TTL", "not-a-number") + assert settings.get_auto_unload_idle_seconds() == 0 + monkeypatch.delenv("UNSLOTH_MODEL_IDLE_TTL", raising = False) + assert settings.get_auto_unload_idle_seconds() == 0 + + +# ── codex/gemini round: standalone-idle reload, path-as-id, embeddings input, retrieve id ── + + +def test_env_idle_standalone_reloads_freed_model_with_auto_switch_off(monkeypatch): + # C3: a standalone UNSLOTH_MODEL_IDLE_TTL (auto-switch OFF) freed the model on + # idle; the next request must restore exactly what was freed even though the + # resolver never runs while auto-switch is off. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = False, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), # would switch if resolver ran + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) # standalone env TTL + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + # Resolver skipped (auto-switch off), so only the stash reload runs: the freed A + # is restored, not the resolves_to target B. + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_no_stash_reload_when_idle_off_and_auto_switch_off(monkeypatch): + # C3 guard: with both auto-switch and idle-unload off the hook is a pure no-op + # and must not resurrect a stashed model (that path only serves the idle feature). + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + _run_hook("org/B-GGUF") + assert rec.calls == [] + + +def test_stash_reload_skipped_while_unsloth_model_active(monkeypatch): + # An Unsloth/Transformers model loaded after an idle-unload leaves the GGUF slot + # empty but is the live model; an unknown /v1 name must NOT resurrect the stale + # GGUF stash (that reload would tear the active Unsloth model down). + from types import SimpleNamespace + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # GGUF slot empty + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + # An Unsloth model is the live backend. + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: SimpleNamespace(active_model_name = "unsloth/Qwen3-8B"), + ) + _run_hook("gpt-4o-mini") + assert rec.calls == [] # stale GGUF not reloaded over the active Unsloth model + + +def test_is_abs_path_id_distinguishes_path_from_repo_id(): + assert resolver._is_abs_path_id("/abs/path/model.gguf") is True + assert resolver._is_abs_path_id("org/Repo-GGUF") is False + assert resolver._is_abs_path_id("Repo") is False + + +def test_advertised_loader_id_prefers_alias_over_abs_path(): + # C1: the ./models and LM Studio scanners report the on-disk path as info.id. + from types import SimpleNamespace + + f = resolver._advertised_loader_id + # An absolute-path id falls back to the first non-path alias. + assert ( + f(SimpleNamespace(id = "/home/me/models/x", model_id = "org/X-GGUF", display_name = "X")) + == "org/X-GGUF" + ) + # No alias available: strip the path to a public id so a host path is never advertised. + assert ( + f( + SimpleNamespace( + id = "/home/me/models/Qwen3-8B-Q4_K_M.gguf", model_id = None, display_name = None + ) + ) + == "Qwen3-8B-Q4_K_M" + ) + # A normal repo id is advertised as-is. + assert ( + f(SimpleNamespace(id = "org/X-GGUF", model_id = "org/X-GGUF", display_name = "X")) == "org/X-GGUF" + ) + + +def test_index_advertises_alias_not_filesystem_path(tmp_path, monkeypatch): + # C1 end-to-end: a scanner that reports the path as the id must not advertise the + # host path in /v1/models, yet the model stays resolvable by that path too. + from types import SimpleNamespace + import routes.models as models_route + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + info = SimpleNamespace( + id = str(gguf), # scanner uses the on-disk path as the id + path = str(gguf), + model_id = "org/Repo-GGUF", + display_name = "Repo", + ) + monkeypatch.setattr(models_route, "_scan_models_dir", lambda *a, **k: [info]) + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + resolver._scan = (0.0, {}) + + # The advertised id is the alias, never the absolute path. + advertised = sorted({entry.loader_id for entry in resolver._index().values()}) + assert advertised == ["org/Repo-GGUF"] + # But the model is still resolvable by its on-disk path (an indexed alias). + resolver._scan = (0.0, {}) + assert resolver.resolve_local_gguf(str(gguf)) is not None + + +def test_build_index_survives_a_failing_scanner(tmp_path, monkeypatch): + # gemini: one bad scanner (e.g. a permission error on ./models) must drop only + # that source, not abort the whole index and lose what the others found. + from types import SimpleNamespace + import routes.models as models_route + import utils.paths as paths + + def _boom(*a, **k): + raise OSError("permission denied") + + lm_info = SimpleNamespace( + id = "org/Repo-GGUF", path = "/lm/Repo", model_id = "org/Repo-GGUF", display_name = "Repo" + ) + monkeypatch.setattr(models_route, "_scan_models_dir", _boom) # ./models blows up + monkeypatch.setattr(models_route, "_scan_hf_cache", lambda *a, **k: []) + monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: tmp_path) + monkeypatch.setattr(models_route, "_is_hidden_model", lambda *a, **k: False) + monkeypatch.setattr(models_route, "_scan_lmstudio_dir", lambda *a, **k: [lm_info]) + monkeypatch.setattr(paths, "legacy_hf_cache_dir", lambda: None) + monkeypatch.setattr(paths, "hf_default_cache_dir", lambda: None) + monkeypatch.setattr(paths, "lmstudio_model_dirs", lambda: [tmp_path]) + # The on-disk GGUF check is covered elsewhere; here a found info becomes an entry. + monkeypatch.setattr( + resolver, + "_local_gguf_entry", + lambda loader_id, info: resolver._LocalGgufEntry(loader_id, "/lm/Repo", ()), + ) + resolver._scan = (0.0, {}) + index = resolver._build_index() + assert any(e.loader_id == "org/Repo-GGUF" for e in index.values()) + + +def test_info_has_local_gguf_reads_files_not_model_format(tmp_path): + # Codex: HF-cache GGUF snapshots leave model_format unset, so /v1/models must + # decide GGUF-ness from the on-disk files. A standalone .gguf (no model_format) + # is servable; a safetensors-only dir is not. + from types import SimpleNamespace + + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(gguf), path = str(gguf))) is True + + st = tmp_path / "safetensors_model" + st.mkdir() + (st / "model.safetensors").write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(st), path = str(st))) is False + + +def test_info_has_local_gguf_excludes_ollama_links(tmp_path): + # Codex P2: Ollama entries come from a scanner _build_index skips, so their + # advertised ids never resolve; the catalog must not report them as servable. + from types import SimpleNamespace + + links = tmp_path / ".studio_links" + links.mkdir() + ollama_gguf = links / "model-Q4_K_M.gguf" + ollama_gguf.write_bytes(b"x" * 32) + assert ( + resolver.info_has_local_gguf(SimpleNamespace(id = "ollama/foo:latest", path = str(ollama_gguf))) + is False + ) + # The same GGUF outside an ollama-link dir is still servable. + plain = tmp_path / "model-Q4_K_M.gguf" + plain.write_bytes(b"x" * 32) + assert resolver.info_has_local_gguf(SimpleNamespace(id = str(plain), path = str(plain))) is True + + +def test_embeddings_input_present_helper(): + f = inference_route._embeddings_input_present + assert f({"input": "hi"}) is True + assert f({"input": ["a", "b"]}) is True + assert f({"input": [1, 2, 3]}) is True + assert f({}) is False + assert f({"input": ""}) is False + assert f({"input": []}) is False + + +def test_embeddings_rejects_missing_input_before_switch(monkeypatch): + # C2: with auto-switch on, an embeddings request carrying no input must 400 + # before the hook, so an invalid request never swaps the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") # loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings(_json_body_request({"model": "org/B-GGUF"}), "tester") + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no model switch happened + + +def test_retrieve_model_tolerates_non_string_id(monkeypatch): + # G2: a model object with a non-string id (defensive) must be skipped rather + # than crashing the .lower() compare; a valid id is still found, unknown 404s. + from fastapi import HTTPException + + async def _objs(): + return [{"id": 123, "object": "model"}, {"id": "org/B-GGUF", "object": "model"}] + + monkeypatch.setattr(inference_route, "_openai_model_objects", lambda: []) # nothing loaded + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _objs) + obj = asyncio.run(inference_route.openai_retrieve_model("org/B-GGUF", "tester")) + assert obj["id"] == "org/B-GGUF" + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_retrieve_model("123", "tester")) + assert exc.value.status_code == 404 + + +def test_retrieve_model_resolves_raw_path_to_advertised_id(monkeypatch): + # Codex P2: a client caching the legacy absolute .gguf path must still retrieve + # a loaded auto-switch model. Its /v1/models entry is keyed by the advertised + # repo id (identifier = snapshot path), so the raw-path fallback must map the raw + # id to that advertised id, not public_model_id(path), or a loaded model 404s. + from types import SimpleNamespace + + raw_path = "/cache/models--org--B-GGUF/snapshots/abc/model.gguf" + llama = SimpleNamespace( + is_loaded = True, model_identifier = raw_path, _openai_advertised_id = "org/B-GGUF" + ) + infer = SimpleNamespace(active_model_name = None) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: infer) + monkeypatch.setattr( + inference_route, + "_openai_model_objects", + lambda: [{"id": "org/B-GGUF", "object": "model"}], + ) + + async def _empty(): + return [] + + monkeypatch.setattr(inference_route, "_openai_catalog_objects", _empty) + obj = asyncio.run(inference_route.openai_retrieve_model(raw_path, "tester")) + assert obj["id"] == "org/B-GGUF" and obj["loaded"] is True + + +def test_chat_streaming_n_gt_1_rejected_before_switch(monkeypatch): + # Codex P2: only the non-streaming GGUF path returns multiple choices, so + # stream=true + n>1 is invalid on every local serving path. Both fields are + # known pre-switch, so it must 400 before the switch rather than loading model B. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", stream = True, n = 2) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_resolver_cache_stamped_after_slow_build(monkeypatch): + # Codex P2: the cache must be stamped AFTER _build_index. A scan slower than the + # TTL would otherwise store an already-expired cache and rebuild every request. + import core.inference.local_model_resolver as r + + clock = {"t": 1000.0} + monkeypatch.setattr(r.time, "monotonic", lambda: clock["t"]) + calls = {"n": 0} + + def _slow_build(): + calls["n"] += 1 + clock["t"] += r._CACHE_TTL_S + 10.0 # the scan itself outlasts the TTL + return {} + + monkeypatch.setattr(r, "_build_index", _slow_build) + r._scan = (0.0, {}) + r._index() # builds once, stamps post-scan + r._index() # immediately after: must reuse the cache, not rebuild + assert calls["n"] == 1 + + +def test_keepwarm_does_not_stamp_activity_on_401(monkeypatch): + # Codex P2: the keep-warm middleware runs before auth, so a 401 must decrement + # the in-flight count without stamping activity, or unauthenticated probes would + # keep the model warm and block idle-unload. + import core.inference.llama_keepwarm as kw + + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_pending", 0) + monkeypatch.setattr(kw, "_last_active", 100.0) + + async def _recv(): + return {"type": "http.request"} + + async def _run(status_code): + async def _app(scope, receive, send): + await send({"type": "http.response.start", "status": status_code, "headers": []}) + await send({"type": "http.response.body", "body": b"x", "more_body": False}) + + sent = [] + + async def _send(m): + sent.append(m) + + mw = kw.LlamaKeepWarmMiddleware(_app) + await mw({"type": "http", "method": "POST", "path": "/v1/chat/completions"}, _recv, _send) + + asyncio.run(_run(401)) + assert kw._inflight == 0 # balanced (start then untracked end) + assert kw._last_active == 100.0 # activity NOT stamped for an auth failure + # A served (200) request still stamps activity. + asyncio.run(_run(200)) + assert kw._inflight == 0 + assert kw._last_active != 100.0 + + +# ── 10-reviewer round: automatic-load validation asymmetry, audio, preview, idle timer ── + + +def _stash(monkeypatch, *, idle = 600): + """Common setup for the standalone-idle reload paths: feature off, idle TTL on, + an idle-freed model in the stash, nothing loaded, no in-flight requests.""" + from core.inference import llama_keepwarm as kw + + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: idle) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + + +def test_completions_prompt_present_helper(): + f = inference_route._completions_prompt_present + assert f({"prompt": "hi"}) is True + assert f({"prompt": ["a", "b"]}) is True + assert f({}) is False + assert f({"prompt": ""}) is False + assert f({"prompt": []}) is False + + +def test_completions_rejects_missing_prompt_before_switch(monkeypatch): + # #1: /v1/completions had no prompt pre-check, so a malformed request naming a + # different downloaded GGUF loaded it before failing. Now it 400s first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF"}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_chat_system_only_rejected_before_idle_reload(monkeypatch): + # #4: the chat pre-load guard only checked auto-switch; a standalone idle TTL + # could still reload a system-only chat before the 400. Now it 400s first. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "system", "content": "sys"}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_embeddings_missing_input_rejected_before_idle_reload(monkeypatch): + # #5: same gap on /v1/embeddings; the missing-input 400 must fire under a + # standalone idle TTL too, not only when auto-switch is on. + from fastapi import HTTPException + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_embeddings(_json_body_request({"model": "x"}), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # no reload before rejection + + +def test_messages_does_not_503_before_reload_hook_when_idle_on(monkeypatch): + # #3: /v1/messages 503'd before the reload hook when auto-switch was off, so a + # standalone idle TTL could never restore the freed model. The early 503 now + # defers to any automatic-load trigger, so the reload hook runs. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + # The handler proceeds past the hook to real generation (no llama-server here), + # so tolerate the downstream failure; the reload having run is the assertion. + try: + asyncio.run( + inference_route.anthropic_messages( + _anthropic_payload(max_tokens = 16), object(), "tester" + ) + ) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_messages_503_gated_on_automatic_load_predicate(): + # Lock the #3 fix at the source: the early 503 must check the shared predicate. + import inspect + src = inspect.getsource(inference_route.anthropic_messages) + assert "_automatic_model_load_may_run" in src + + +def test_raw_body_without_model_reloads_freed_model(monkeypatch): + # #6: a raw completions/embeddings body that omits `model` passed None, which + # skipped the idle-stash reload and 503'd. A non-empty sentinel now lets the + # reload run while still resolving as unknown. + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + assert rec.calls[0].gguf_variant == "Q4_K_M" + + +def test_audio_generate_reloads_idle_freed_model(monkeypatch): + # #2: /audio/generate is keep-warm-tracked but had no reload hook, so an + # idle-freed audio GGUF stayed unloaded. The hook now restores it. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = [{"role": "user", "content": "say hi"}]) + # Falls through to the non-audio backend path (no real model) after the reload; + # tolerate that downstream failure, the reload having run is the assertion. + try: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + except Exception: + pass + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def test_audio_generate_does_not_reload_on_invalid_request(monkeypatch): + # The audio reload hook must run after message validation, so an empty request + # never triggers a reload. + from fastapi import HTTPException + from models.inference import ChatCompletionRequest + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + _stash(monkeypatch) + payload = ChatCompletionRequest(model = "x", messages = []) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_preview_scope_disables_auto_switch(monkeypatch): + # #7: the public preview route delegates to the chat handler; a caller-supplied + # model must not switch away from the pinned checkpoint. The scope opt-out flag + # makes the hook a no-op. + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + + class _Req: + def __init__(self): + self.scope = {} + + req = _Req() + inference_route.disable_openai_auto_switch_for_request(req.scope) + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req, "tester")) + assert rec.calls == [] # preview opt-out suppressed the switch + + # Control: a fresh request without the flag would switch. + req2 = _Req() + asyncio.run(inference_route._maybe_auto_switch_model("org/B-GGUF", req2, "tester")) + assert len(rec.calls) == 1 + + +def test_preview_chat_is_tracked_as_inference_path(): + # #8: long preview streams use the same backend; the keep-warm middleware must + # count them so the idle loop can't unload mid-response. + from core.inference.llama_keepwarm import _is_inference_path + + assert _is_inference_path("/p/my-run/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/ckpt-100/v1/chat/completions") is True + assert _is_inference_path("/p/my-run/v1/models") is False + + +def test_untrack_does_not_reset_idle_timer(): + # #9: external-provider traffic was keeping the local GGUF warm forever because + # untrack stamped _last_active. It must decrement in-flight without restamping. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 1 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + scope = {"type": "http"} + kw.untrack_current_request(scope) + assert kw._inflight == 0 + assert kw._last_active == before # idle timer not reset by an untracked request + kw._inflight = 0 + + +def test_note_start_does_not_reset_idle_timer(): + # The start stamp was removed so an external request that is later untracked + # cannot reset the timer at start either; in-flight count still protects it. + import time + from core.inference import llama_keepwarm as kw + + kw._inflight = 0 + kw._pending = 0 + kw._last_active = time.monotonic() - 3600 + before = kw._last_active + kw._note_start() + try: + assert kw._inflight == 1 + assert kw._last_active == before # start no longer stamps activity + assert kw._is_idle(1.0) is False # but in-flight still blocks unload + finally: + kw._note_end() # restores _last_active stamp on completion + + +# ── codex review (merge round): reload-only sentinel, Anthropic tool validation ── + + +def test_omitted_model_does_not_resolve_to_a_named_gguf(monkeypatch): + # Codex P2: a raw-body request that omits `model` must never run the resolver, + # so a downloaded GGUF literally named "default" can't be switched to. The + # resolver here would switch to B if it ran; it must not. + backend = _FakeBackend("org/A-GGUF") # a model is already loaded + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + body = asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert body == {"prompt": "hi"} + assert rec.calls == [] # resolver skipped (would have switched to B otherwise) + + +def test_omitted_model_still_reloads_idle_freed_model(monkeypatch): + # The reload-only sentinel must still restore an idle-freed model (the round-9 + # behavior), it just never runs the resolver. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) # idle-unload emptied the slot + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + asyncio.run( + inference_route._auto_switch_from_request_body( + _json_body_request({"prompt": "hi"}), "tester" + ) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" + + +def _anthropic_payload_with_tools(tools, max_tokens = 16): + from models.inference import AnthropicMessagesRequest, AnthropicMessage + return AnthropicMessagesRequest( + model = "org/B-GGUF", + max_tokens = max_tokens, + messages = [AnthropicMessage(role = "user", content = "hi")], + tools = tools, + ) + + +def test_anthropic_invalid_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed client tool (no input_schema, no server-tool type) must + # 400 before the auto-switch hook, so an invalid request never evicts the model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # missing input_schema + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_anthropic_validates_tools_before_auto_switch(): + # Lock the order at the source: tool-shape validation precedes the hook, for + # both /messages and /messages/count_tokens (shared helper). + import inspect + for fn in (inference_route.anthropic_messages, inference_route.anthropic_count_tokens): + src = inspect.getsource(fn) + assert src.index("_validate_anthropic_client_tools") < src.index("_maybe_auto_switch_model") + + +def test_anthropic_mixed_tools_rejected_before_switch(monkeypatch): + # Codex P2: combining an Anthropic server tool (type) with a custom client tool + # (input_schema) is unsupported and must 400 before the switch, so the request + # can't evict the loaded model only to be rejected after the load. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools( + [ + {"type": "web_search_20250305"}, # server tool + {"name": "my_func", "input_schema": {"type": "object"}}, # client tool + ] + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_messages(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +# ── codex review (round 2): schema-default model, Responses tool validation ── + + +def _chat_msg(text = "hi"): + from models.inference import ChatMessage + return ChatMessage(role = "user", content = text) + + +def _responses_payload(*, tools = None, set_model = True): + from models.inference import ResponsesRequest + + kwargs = dict(input = "hi") + if set_model: + kwargs["model"] = "org/B-GGUF" + if tools is not None: + kwargs["tools"] = tools + return ResponsesRequest(**kwargs) + + +def test_switch_model_for_payload_only_switches_when_explicit(): + # Codex P2: an omitted `model` (pydantic fills "default") must be reload-only; + # an explicitly set model -- including a literal "default" -- is honored. + from models.inference import ChatCompletionRequest + + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(omitted) == inference_route._RELOAD_ONLY_MODEL + explicit_default = ChatCompletionRequest(model = "default", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit_default) == "default" + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + assert inference_route._switch_model_for_payload(explicit) == "org/B-GGUF" + + +def test_omitted_schema_model_skips_resolver(monkeypatch): + # End to end: a schema request omitting `model` must not run the resolver, so a + # GGUF named "default" is never swapped to; an explicit model still switches. + from models.inference import ChatCompletionRequest + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + omitted = ChatCompletionRequest(messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(omitted), object(), "tester" + ) + ) + assert rec.calls == [] # resolver skipped + explicit = ChatCompletionRequest(model = "org/B-GGUF", messages = [_chat_msg()]) + asyncio.run( + inference_route._maybe_auto_switch_model( + inference_route._switch_model_for_payload(explicit), object(), "tester" + ) + ) + assert len(rec.calls) == 1 # explicit model still switches + + +def test_build_chat_request_propagates_omitted_model(): + # _build_chat_request must not turn an omitted Responses model into an explicit + # "default", or the non-streaming chat re-check would switch on it. + omitted = _responses_payload(set_model = False) + chat_req = inference_route._build_chat_request(omitted, [_chat_msg()], stream = False) + assert "model" not in chat_req.model_fields_set + explicit = _responses_payload(set_model = True) + chat_req2 = inference_route._build_chat_request(explicit, [_chat_msg()], stream = False) + assert "model" in chat_req2.model_fields_set + + +def test_responses_invalid_function_tool_rejected_before_switch(monkeypatch): + # Codex P2: a malformed function tool (no name) must 400 before the hook, so an + # invalid /v1/responses request never switches or evicts the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _responses_payload(tools = [{"type": "function", "parameters": {}}]) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the model load + + +def test_responses_valid_and_builtin_tools_pass_validation(monkeypatch): + # A well-formed function tool and a built-in (non-function) tool must pass the + # pre-switch check. Stub the hook so the test stops right after validation. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _responses_payload( + tools = [{"type": "function", "name": "ok", "parameters": {}}, {"type": "web_search"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + + +def test_responses_validates_tools_before_auto_switch(): + # Lock the order at the source: tool validation precedes the switch hook. + import inspect + src = inspect.getsource(inference_route.openai_responses) + assert src.index("each function tool must have a 'name'") < src.index( + "_maybe_auto_switch_model" + ) + + +def test_responses_forcing_tool_choice_without_name_rejected_before_switch(monkeypatch): + # Codex P2: a forcing-function tool_choice with no name (Responses shape + # {"type": "function"}) must 400 before the switch, so the streaming path can't + # forward a bad choice and an invalid request can't evict the model. + from fastapi import HTTPException + from models.inference import ResponsesRequest + + async def _boom(*a, **k): + raise AssertionError("must not switch on an invalid tool_choice") + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = ResponsesRequest(model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function"}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_responses(payload, object(), "tester")) + assert exc.value.status_code == 400 + # A named forcing choice is accepted (reaches the switch, which is mocked to raise). + ok = ResponsesRequest( + model = "org/B-GGUF", input = "hi", tool_choice = {"type": "function", "name": "f"} + ) + with pytest.raises(AssertionError): + asyncio.run(inference_route.openai_responses(ok, object(), "tester")) + + +# ── codex review (round 3): process-wide swap gate across event loops ── + + +def test_swap_acquires_process_gate_before_load(): + # Lock in the structure: the process-wide gate is acquired before the load and + # always released, so a cross-loop swap can't reach _load_model_impl unguarded. + import inspect + + src = inspect.getsource(inference_route._maybe_auto_switch_model) + assert src.index("_acquire_swap_gate") < src.index("_load_model_impl") + assert "_auto_switch_process_lock.release()" in src + + +# ── codex review (round 4): validate modality + tool-confirmation before switch ── + + +def _chat_request(**kw): + from models.inference import ChatCompletionRequest, ChatMessage + kw.setdefault("messages", [ChatMessage(role = "user", content = "hi")]) + return ChatCompletionRequest(**kw) + + +def test_chat_confirm_without_stream_rejected_before_switch(monkeypatch): + # Codex P2: confirm_tool_calls=true + stream=false + local tools is an invalid + # shape; it must 400 before the switch hook so it can't evict the resident model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", enable_tools = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_confirm_with_bypass_permissions_reaches_hook(monkeypatch): + # bypass_permissions suppresses the confirm gate, so the pre-check must not fire; + # the request should reach the switch hook (stubbed here to a sentinel). + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", + enable_tools = True, + confirm_tool_calls = True, + stream = False, + bypass_permissions = True, + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_chat_audio_input_guards_target_before_switch(monkeypatch): + # Codex P2: a chat request carrying audio_base64 must guard the target before the + # switch -- audio rides the same companion mmproj as vision -- so a text-only + # target can't be loaded and evict the working audio model. Assert the handler + # flags require_vision so the hook's multimodal probe runs. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = "AAAA") + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_completions_rejects_object_prompt_before_switch(monkeypatch): + # Codex P2: an object prompt like {"prompt": {}} is a deterministic client error + # (only a string or array is valid). It must 400 before the switch so a bad shape + # can't load the named GGUF only to be rejected by llama-server after eviction. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_completions( + _json_body_request({"model": "org/B-GGUF", "prompt": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # no switch before rejection + + +def test_embeddings_rejects_object_input_before_switch(monkeypatch): + # Codex P2: an object input like {"input": {}} is a deterministic client error + # (only a string or array is valid); reject before the switch, like completions. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route.openai_embeddings( + _json_body_request({"model": "org/B-GGUF", "input": {}}), "tester" + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_oversized_audio_rejected_before_switch(monkeypatch): + # Codex P2: the audio size cap is a cheap, target-independent length check, so an + # oversized upload must 413 before the switch rather than loading a GGUF first. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + big = "A" * (inference_route._MAX_AUDIO_B64_CHARS + 1) + payload = _chat_request(model = "org/B-GGUF", audio_base64 = big) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 413 + assert rec.calls == [] + + +def test_chat_confirm_without_stream_mcp_rejected_before_switch(monkeypatch): + # Codex P2: mcp_enabled opens the local tool loop on its own, so confirm+no-stream + # +mcp is the same invalid shape as confirm+no-stream+tools and must 400 before + # the switch. The old guard only checked explicit tool fields and missed it. + import state.tool_policy as _tp + from fastapi import HTTPException + + monkeypatch.setattr(_tp, "get_tool_policy", lambda: None) # no CLI --disable-tools + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request( + model = "org/B-GGUF", mcp_enabled = True, confirm_tool_calls = True, stream = False + ) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_require_vision_rejects_text_target_before_switch(monkeypatch): + # Codex P2: an image request naming a different text-only GGUF must 400 before + # the swap, so the resident vision model is not evicted for a rejected request. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: False) + with pytest.raises(HTTPException) as exc: + asyncio.run( + inference_route._maybe_auto_switch_model( + "org/B-GGUF", object(), "t", require_vision = True + ) + ) + assert exc.value.status_code == 400 + assert rec.calls == [] # rejected before the load + + +def test_require_vision_allows_vision_target(monkeypatch): + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/local/B.gguf", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + monkeypatch.setattr(inference_route, "_target_is_vision", lambda _p: True) + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 # vision target still switches + + +def test_require_vision_ignores_reload_stash(monkeypatch): + # The reload-stash path restores the model the request was already using; the + # modality check applies only to an explicit resolver target, not a restore. + from core.inference import llama_keepwarm as kw + + backend = _FakeBackend(None) + rec = _LoadRecorder(backend) + _wire(monkeypatch, enabled = False, resolves_to = None, backend = backend, recorder = rec) + monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 600) + monkeypatch.setattr(kw, "_inflight", 0) + monkeypatch.setattr(kw, "_last_unloaded_model", ("/cache/snap/A", "Q4_K_M", "org/A-GGUF")) + monkeypatch.setattr( + inference_route, "_target_is_vision", lambda _p: False + ) # would reject if used + asyncio.run( + inference_route._maybe_auto_switch_model("org/B-GGUF", object(), "t", require_vision = True) + ) + assert len(rec.calls) == 1 + assert rec.calls[0].model_path == "/cache/snap/A" # restored despite require_vision + + +def test_chat_validates_confirm_and_modality_before_switch(): + # Lock the order at the source: confirm-shape rejection precedes the hook, and + # the hook rejects a non-vision target before the load. + import inspect + + src = inspect.getsource(inference_route.openai_chat_completions) + assert src.index("confirm_tool_calls requires stream=true") < src.index( + "_maybe_auto_switch_model" + ) + assert "require_vision" in src + hook = inspect.getsource(inference_route._maybe_auto_switch_model) + assert hook.index("require_vision") < hook.index("_load_model_impl") + assert "does not support the image or audio input" in hook + + +def test_messages_have_image_helper(): + from models.inference import ChatMessage, ImageContentPart, ImageUrl, TextContentPart + + f = inference_route._messages_have_image + text_only = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage(role = "user", content = [TextContentPart(type = "text", text = "hi")]), + ] + assert f(text_only) is False + img = ImageContentPart(type = "image_url", image_url = ImageUrl(url = "data:image/png;base64,AAAA")) + assert f([ChatMessage(role = "user", content = [img])]) is True + + +def test_anthropic_request_has_image_helper(): + from types import SimpleNamespace + + f = inference_route._anthropic_request_has_image + text = SimpleNamespace(messages = [SimpleNamespace(content = "hi")]) + assert f(text) is False + text_block = SimpleNamespace( + messages = [SimpleNamespace(content = [{"type": "text", "text": "hi"}])] + ) + assert f(text_block) is False + dict_img = SimpleNamespace(messages = [SimpleNamespace(content = [{"type": "image"}])]) + assert f(dict_img) is True + typed_img = SimpleNamespace(messages = [SimpleNamespace(content = [SimpleNamespace(type = "image")])]) + assert f(typed_img) is True + + +def test_responses_and_anthropic_wire_require_vision_from_images(): + # P2: the modality guard must fire on /v1/responses and /v1/messages too, so an + # image request can't evict a vision model for a text-only target. Lock the wiring + # at the source: each hook derives require_vision from the request's images. + import inspect + + responses_src = inspect.getsource(inference_route.openai_responses) + assert "require_vision = _messages_have_image(" in responses_src + anthropic_src = inspect.getsource(inference_route.anthropic_messages) + assert "require_vision = _anthropic_request_has_image(" in anthropic_src + # /messages/count_tokens shares the /messages translation, so it needs the same + # guard: an image count must not evict a vision model for a text-only target. + count_src = inspect.getsource(inference_route.anthropic_count_tokens) + assert "require_vision = _anthropic_request_has_image(" in count_src + + +# ── codex review (round 5): count_tokens tools, tool_choice, process-wide gate ── + + +def test_count_tokens_rejects_malformed_tool_before_switch(monkeypatch): + # Codex P2: /v1/messages/count_tokens must reject a malformed tool before the + # switch, like /messages, so a count request can't evict the loaded model. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("/p/B", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _anthropic_payload_with_tools([{"name": "broken"}]) # no input_schema/type + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_count_tokens_forwards_vision_guard_to_switch(monkeypatch): + # Codex P2: an image /v1/messages/count_tokens naming a text-only GGUF must + # carry the same require_vision guard as /messages, so it can't evict a loaded + # vision model for a swap that can't serve the request. + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["require_vision"] = require_vision + raise _Reached() + + monkeypatch.setattr(inference_route, "_anthropic_request_has_image", lambda p: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = _anthropic_payload_with_tools(None) # no tools -> tool validation passes + with pytest.raises(_Reached): + asyncio.run(inference_route.anthropic_count_tokens(payload, object(), "tester")) + assert captured["require_vision"] is True + + +def test_audio_generate_is_reload_only(monkeypatch): + # Codex P2: /audio/generate must not switch to a client-named GGUF. A local + # GGUF's audio-input capability is not a cheap pre-load probe (the mmproj signal + # can't tell an audio projector from a vision one), so resolving the client model + # could evict the working audio model for a target that then fails the audio + # check. Only the idle-stash restore runs: the hook gets the reload-only sentinel. + from models.inference import ChatCompletionRequest + + class _Reached(Exception): + pass + + captured = {} + + async def _capture( + model, + request, + subject, + *, + require_vision = False, + ): + captured["model"] = model + raise _Reached() + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _capture) + payload = ChatCompletionRequest( + model = "org/B-GGUF", messages = [{"role": "user", "content": "say hi"}] + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.generate_audio(payload, object(), "tester")) + assert captured["model"] == inference_route._RELOAD_ONLY_MODEL + + +def test_note_model_unloaded_clears_reload_stash(monkeypatch): + # Codex P2: a deliberate unload must drop the idle reload stash so the next /v1 + # request can't resurrect the just-unloaded model. (The idle loop unloads via the + # backend directly, so clearing on the route never fights keep-warm.) + import core.inference.llama_keepwarm as kw + + kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M")) + assert kw.get_last_unloaded_model() == ("org/A-GGUF", "Q4_K_M") + kw.note_model_unloaded() + assert kw.get_last_unloaded_model() is None + + +def test_unload_route_clears_reload_stash(monkeypatch): + # The /unload route must clear the stash on both the GGUF and non-GGUF branches. + import inspect + src = inspect.getsource(inference_route.unload_model) + assert src.count("note_model_unloaded()") >= 2 + + +def test_non_gguf_load_clears_reload_stash(): + # A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF + # branch, so it never lingers until the idle poll (or forever, idle-unload off). + import inspect + src = inspect.getsource(inference_route._load_model_impl) + assert src.count("note_model_loaded()") >= 2 + + +def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch): + # Codex P2: a forcing object with no function name must 400 before the switch. + from fastapi import HTTPException + + backend = _FakeBackend("org/A-GGUF") + rec = _LoadRecorder(backend) + _wire( + monkeypatch, + enabled = True, + resolves_to = ("org/B-GGUF", "Q8_0", "org/B-GGUF"), + backend = backend, + recorder = rec, + ) + payload = _chat_request(model = "org/B-GGUF", tool_choice = {"type": "function", "function": {}}) + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + assert exc.value.status_code == 400 + assert rec.calls == [] + + +def test_chat_valid_tool_choice_reaches_hook(monkeypatch): + # A well-formed forcing object must pass the pre-check and reach the hook. + class _Reached(Exception): + pass + + async def _boom(*a, **k): + raise _Reached() + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _boom) + payload = _chat_request( + model = "org/B-GGUF", tool_choice = {"type": "function", "function": {"name": "ok"}} + ) + with pytest.raises(_Reached): + asyncio.run(inference_route.openai_chat_completions(payload, object(), "tester")) + + +def test_lifecycle_gate_serializes_across_loops(): + # Codex P2: the lifecycle gate must be process-wide so a swap on one loop blocks + # inference starting on another. Two loops must never hold the gate at once. + import threading + from core.inference import llama_keepwarm as kw + + state = {"cur": 0, "max": 0} + slock = threading.Lock() + + async def _use(): + async with kw._unload_gate(): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.05) + with slock: + state["cur"] -= 1 + + barrier = threading.Barrier(2) + + def _run(): + barrier.wait() + asyncio.run(_use()) + + threads = [threading.Thread(target = _run) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join() + assert state["max"] == 1 # never held on two loops at once + + +def test_auto_switch_serializes_across_event_loops(monkeypatch): + # Codex P2: the per-loop asyncio lock can't serialize two swaps on different + # event loops in one process. The process-wide gate must, so the two slow loads + # never overlap on the single model slot. + import threading + + backend = _FakeBackend("org/A-GGUF") + state = {"cur": 0, "max": 0} + loaded: list = [] + slock = threading.Lock() + + async def _slow_load( + request, + fastapi_request, + current_subject = None, + ): + with slock: + state["cur"] += 1 + state["max"] = max(state["max"], state["cur"]) + await asyncio.sleep(0.1) # widen the window so an unguarded race would overlap + with slock: + state["cur"] -= 1 + loaded.append(request.model_path) + backend.model_identifier = request.model_path + backend.is_loaded = True + backend._openai_advertised_id = None + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + monkeypatch.setattr(resolver, "resolve_local_gguf", lambda m: (m, "Q8_0", m)) + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inference_route, "_load_model_impl", _slow_load) + monkeypatch.setattr(inference_route, "_auto_switch_waiters", {}) + monkeypatch.setattr(inference_route, "_auto_switch_request_waiters", {}) + + barrier = threading.Barrier(2) + + def _run(model): + barrier.wait() # release both threads together so they truly race + asyncio.run(inference_route._maybe_auto_switch_model(model, object(), "t")) + + threads = [ + threading.Thread(target = _run, args = ("org/B-GGUF",)), + threading.Thread(target = _run, args = ("org/C-GGUF",)), + ] + for t in threads: + t.start() + for t in threads: + t.join() + + assert state["max"] == 1 # the gate serialized the two cross-loop swaps + assert sorted(loaded) == ["org/B-GGUF", "org/C-GGUF"] # both still swapped + + +def test_acquire_swap_gate_is_cancellation_safe(): + # A waiter cancelled while waiting for the gate (client disconnect mid-swap) + # must not leak it: after the holder releases, a fresh acquire still succeeds. + # The to_thread(acquire) approach would leak here -- its worker thread keeps + # acquiring after cancel, so the gate is taken but never released. + async def main(): + await inference_route._acquire_swap_gate() # this loop holds the gate + try: + + async def waiter(): + await inference_route._acquire_swap_gate() + + t = asyncio.create_task(waiter()) + await asyncio.sleep(0.05) # let it spin waiting on the held gate + t.cancel() + with pytest.raises(asyncio.CancelledError): + await t + finally: + inference_route._auto_switch_process_lock.release() + # Gate is free again (the cancelled waiter never acquired it). + await asyncio.wait_for(inference_route._acquire_swap_gate(), timeout = 1) + inference_route._auto_switch_process_lock.release() + + asyncio.run(asyncio.wait_for(main(), timeout = 5)) diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py index f9baf20a66..552f122ebb 100644 --- a/studio/backend/tests/test_openai_catalog.py +++ b/studio/backend/tests/test_openai_catalog.py @@ -13,6 +13,7 @@ if str(_BACKEND) not in sys.path: sys.path.insert(0, str(_BACKEND)) import routes.inference as inf # noqa: E402 +from core.inference import local_model_resolver as resolver # noqa: E402 class _Info: @@ -21,10 +22,12 @@ class _Info: id, display_name, model_id = None, + is_gguf = True, ): self.id = id self.display_name = display_name self.model_id = model_id + self.is_gguf = is_gguf # drives the files-based GGUF check in the test class _FakeLlama: @@ -53,10 +56,16 @@ def test_catalog_lists_loaded_and_available(monkeypatch): return [ _Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup _Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded - _Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id + # HF-cache GGUF: model_format is unset for these, so a files-based check + # (not model_format) must still list it. + _Info("models--org--Foo", "Foo", model_id = "org/Foo"), + # Non-GGUF (safetensors) can't be served via /v1: must NOT be advertised. + _Info("/data/models/Mistral-7B", "Mistral-7B", is_gguf = False), ] monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog) + # GGUF-ness is read from the on-disk files; drive it off each info's flag here. + monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf) data = asyncio.run(inf._openai_catalog_objects()) ids = {m["id"]: m for m in data} @@ -64,9 +73,12 @@ def test_catalog_lists_loaded_and_available(monkeypatch): # Loaded model is present, marked loaded, and keeps context fields. assert ids["Qwen3-Q4"]["loaded"] is True assert ids["Qwen3-Q4"]["context_length"] == 4096 - # Available-but-not-loaded models are listed too. + # Available-but-not-loaded GGUF models are listed too. assert ids["Llama-8B-Q8"]["loaded"] is False + # The HF-cache GGUF is listed despite model_format being unset. assert ids["org/Foo"]["loaded"] is False + # The non-GGUF model is filtered out (/v1 can never serve it). + assert "Mistral-7B" not in ids # The loaded gguf and the on-disk copy collapse to one clean id. assert [m["id"] for m in data].count("Qwen3-Q4") == 1 # No absolute paths or .gguf suffixes leak anywhere. @@ -76,6 +88,20 @@ def test_catalog_lists_loaded_and_available(monkeypatch): assert "/data/" not in blob +def test_catalog_lock_is_per_loop(): + # Codex P2: a module-level asyncio.Lock ties its waiters to the loop that first + # awaited it, so a second event loop awaiting it in a multi-loop process can + # hang. The catalog lock must be per-loop (distinct lock per running loop), and + # the old shared _CATALOG_LOCK must be gone so it can't be reintroduced. + async def _get(): + return inf._catalog_lock() + + a = asyncio.run(_get()) + b = asyncio.run(_get()) # a fresh event loop + assert a is not b + assert not hasattr(inf, "_CATALOG_LOCK") + + def test_empty_and_errored_scans_are_cached(monkeypatch): # Cache validity is keyed on the timestamp, not list contents, so an empty # (fresh install / no local models) or errored scan is still cached for the diff --git a/studio/backend/utils/openai_auto_switch_settings.py b/studio/backend/utils/openai_auto_switch_settings.py new file mode 100644 index 0000000000..1689395f40 --- /dev/null +++ b/studio/backend/utils/openai_auto_switch_settings.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted opt-in controls for OpenAI-compatible model auto-switching. + +Two settings, both off by default so existing API behavior is unchanged: +- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model`` + names a downloaded local GGUF different from the loaded one transparently + loads it before serving (llama-swap-style). Unknown names pass through. +- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is + unloaded after this many idle seconds to free VRAM. + +The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env +var. Unlike the stored setting (which stays gated on auto-switch), the env value +is a standalone default that enables idle-unload even with auto-switch off, for +headless/container deploys; an explicit UI/API value still overrides it. + +Reads are cached for a short window because these are consulted on the +per-request hot path; writes invalidate the cache. +""" + +from __future__ import annotations + +import os +import threading +import time +from typing import Any, Optional + +OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model" +AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds" +MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides" +MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL" + +DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False +DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0 + +_CACHE_TTL_S = 2.0 +_cache_lock = threading.Lock() +_cache: dict[str, tuple[float, Any]] = {} + + +def _coerce_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return None + + +def _coerce_int(value: Any) -> int | None: + try: + return max(0, int(value)) + except (TypeError, ValueError): + return None + + +def _cached_setting(key: str, default: Any) -> Any: + """Read an app setting, memoized for _CACHE_TTL_S to spare the hot path.""" + now = time.monotonic() + with _cache_lock: + hit = _cache.get(key) + if hit is not None and now - hit[0] < _CACHE_TTL_S: + return hit[1] + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(key, None) + except Exception: + stored = None + value = default if stored is None else stored + with _cache_lock: + _cache[key] = (now, value) + return value + + +def _invalidate(key: str) -> None: + with _cache_lock: + _cache.pop(key, None) + + +def get_openai_auto_switch_enabled() -> bool: + parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_SWITCH_SETTING_KEY, None)) + return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED + + +def _stored_idle_seconds() -> Optional[int]: + """The persisted idle TTL as an int, or None when never set.""" + return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None)) + + +def _env_idle_seconds() -> Optional[int]: + """UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid.""" + raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR) + if raw is None or not raw.strip(): + return None + return _coerce_int(raw) + + +def get_stored_auto_unload_idle_seconds() -> int: + """The persisted idle-unload TTL, independent of whether auto-switch is on. + + The settings UI reads this so it can display and round-trip the saved value; + toggling auto-switch off must not erase it. Falls back to the env override so + the UI shows the startup default. The idle loop uses the gated reader below. + """ + stored = _stored_idle_seconds() + if stored is not None: + return stored + env = _env_idle_seconds() + return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS + + +def get_auto_unload_idle_seconds() -> int: + """Effective idle TTL the idle loop runs on (0 = never unload).""" + stored = _stored_idle_seconds() + if stored is not None: + # An explicit UI/API value stays gated on auto-switch: off reports 0 so the + # off state is identical to pre-feature. + return stored if get_openai_auto_switch_enabled() else 0 + # No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that + # enables idle-unload even with auto-switch off (headless/container deploys). + env = _env_idle_seconds() + return env if env is not None else 0 + + +def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]: + """Set both auto-switch flags in one transaction so a settings PUT can't leave + one key updated and the other stale. Both values are coerced before any write, + so an invalid value raises without persisting either.""" + parsed_enabled = _coerce_bool(enabled) + if parsed_enabled is None: + raise ValueError("OpenAI auto-switch must be true or false.") + parsed_idle = _coerce_int(idle_seconds) + if parsed_idle is None: + raise ValueError("Auto-unload idle seconds must be a non-negative integer.") + from storage.studio_db import upsert_app_settings + + upsert_app_settings( + {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle} + ) + _invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY) + _invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY) + return parsed_enabled, parsed_idle + + +def get_model_overrides() -> dict[str, dict]: + """Per-model launch overrides keyed by model id ({llama_extra_args, max_seq_length}).""" + raw = _cached_setting(MODEL_OVERRIDES_SETTING_KEY, None) + return raw if isinstance(raw, dict) else {} + + +def get_model_override(model_id: str) -> dict: + """The launch override applied when auto-switch loads ``model_id`` (or empty).""" + override = get_model_overrides().get(model_id) + return override if isinstance(override, dict) else {} + + +def set_model_override( + model_id: str, + llama_extra_args: Optional[list[str]] = None, + max_seq_length: Optional[int] = None, +) -> dict: + """Upsert one model's launch override; an override with no fields removes it.""" + if not model_id or not model_id.strip(): + raise ValueError("model_id is required.") + entry: dict[str, Any] = {} + if llama_extra_args: + entry["llama_extra_args"] = [str(arg) for arg in llama_extra_args] + if max_seq_length: + entry["max_seq_length"] = max(0, int(max_seq_length)) + + from storage.studio_db import upsert_app_setting_map_entry + + # Atomic per-entry merge so two PUTs for different models can't drop each other. + upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None) + _invalidate(MODEL_OVERRIDES_SETTING_KEY) + return entry diff --git a/studio/frontend/src/features/settings/api/openai-auto-switch.ts b/studio/frontend/src/features/settings/api/openai-auto-switch.ts new file mode 100644 index 0000000000..80ffc084d0 --- /dev/null +++ b/studio/frontend/src/features/settings/api/openai-auto-switch.ts @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type OpenAIAutoSwitchSettings = { + enabled: boolean; + autoUnloadIdleSeconds: number; + defaultEnabled: boolean; + // True when the idle-unload loop will actually unload (e.g. enabled via the + // UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off). + idleUnloadActive: boolean; +}; + +type ApiOpenAIAutoSwitchSettings = { + enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: number; + // biome-ignore lint/style/useNamingConvention: API schema + default_enabled: boolean; + // biome-ignore lint/style/useNamingConvention: API schema + idle_unload_active?: boolean; +}; + +let cachedSettings: OpenAIAutoSwitchSettings | null = null; +let inFlightSettings: Promise | null = null; + +function fromApi( + settings: ApiOpenAIAutoSwitchSettings, +): OpenAIAutoSwitchSettings { + return { + enabled: settings.enabled, + autoUnloadIdleSeconds: settings.auto_unload_idle_seconds, + defaultEnabled: settings.default_enabled, + idleUnloadActive: settings.idle_unload_active ?? false, + }; +} + +async function fetchOpenAIAutoSwitchSettings(): Promise { + const res = await authFetch("/api/settings/openai-auto-switch"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load model auto-switch settings"), + ); + } + return fromApi(await res.json()); +} + +function cacheSettings(settings: OpenAIAutoSwitchSettings) { + cachedSettings = settings; + return settings; +} + +export async function loadOpenAIAutoSwitchSettings() { + if (cachedSettings) { + return cachedSettings; + } + inFlightSettings ??= fetchOpenAIAutoSwitchSettings() + .then(cacheSettings) + .finally(() => { + inFlightSettings = null; + }); + return inFlightSettings; +} + +export async function updateOpenAIAutoSwitchSettings( + enabled: boolean, + autoUnloadIdleSeconds: number, +): Promise { + const res = await authFetch("/api/settings/openai-auto-switch", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + enabled, + // biome-ignore lint/style/useNamingConvention: API schema + auto_unload_idle_seconds: autoUnloadIdleSeconds, + }), + }); + if (!res.ok) { + throw new Error( + await readFastApiError( + res, + "Failed to update model auto-switch settings", + ), + ); + } + return cacheSettings(fromApi(await res.json())); +} diff --git a/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx new file mode 100644 index 0000000000..581c78ebd1 --- /dev/null +++ b/studio/frontend/src/features/settings/components/model-auto-switch-section.tsx @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Switch } from "@/components/ui/switch"; +import { useT } from "@/i18n"; +import { useEffect, useState } from "react"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; +import { SettingsRow } from "./settings-row"; +import { SettingsSection } from "./settings-section"; + +export function ModelAutoSwitchSection() { + const t = useT(); + const [settings, setSettings] = useState( + null, + ); + const [draftIdleSeconds, setDraftIdleSeconds] = useState("0"); + const [error, setError] = useState(null); + const [isSaving, setIsSaving] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((loaded) => { + if (cancelled) return; + setSettings(loaded); + setDraftIdleSeconds(String(loaded.autoUnloadIdleSeconds)); + setError(null); + }) + .catch((loadError) => { + if (cancelled) return; + setError( + loadError instanceof Error + ? loadError.message + : t("settings.general.modelAutoSwitch.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + + // Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null. + const parseIdleSeconds = (): number | null => { + if (!draftIdleSeconds.trim()) { + return null; + } + const parsed = Number(draftIdleSeconds); + return Number.isInteger(parsed) && parsed >= 0 ? parsed : null; + }; + + const persist = async ( + enabled: boolean, + idleSeconds: number, + syncDraft = true, + ) => { + setIsSaving(true); + setError(null); + try { + const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds); + setSettings(saved); + if (syncDraft) { + setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds)); + } + } catch (saveError) { + setError( + saveError instanceof Error + ? saveError.message + : t("settings.general.modelAutoSwitch.saveError"), + ); + } finally { + setIsSaving(false); + } + }; + + // Idle-unload is tied to auto-switch (the freed model reloads via the swap). + // Toggling off preserves the saved seconds rather than zeroing them — the + // backend gates unloading on the enabled flag, so it never unloads while off. + // Enabling commits the drafted value, falling back to the last saved one so + // it can never get stuck. + const handleToggle = (enabled: boolean) => { + const savedIdleSeconds = settings?.autoUnloadIdleSeconds ?? 0; + if (!enabled) { + void persist(false, savedIdleSeconds, false); + return; + } + void persist(true, parseIdleSeconds() ?? savedIdleSeconds); + }; + + const handleSaveIdle = () => { + const idleSeconds = parseIdleSeconds(); + if (idleSeconds === null) { + setError(t("settings.general.modelAutoSwitch.idleError")); + return; + } + void persist(true, idleSeconds); + }; + + return ( + + + + + +
+
+
+ setDraftIdleSeconds(event.target.value)} + className="h-8 w-full pr-8" + /> + + s + +
+ +
+ {error ? ( + + {error} + + ) : settings && !settings.enabled && settings.idleUnloadActive ? ( + + {t("settings.general.modelAutoSwitch.idleActiveViaEnv")} + + ) : settings && !settings.enabled ? ( + + {t("settings.general.modelAutoSwitch.idleNeedsEnable")} + + ) : null} +
+ + + ); +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index 4085aa2ab0..aef2e7ffe6 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -27,6 +27,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useMemo, useState } from "react"; import { Streamdown } from "streamdown"; +import { + type OpenAIAutoSwitchSettings, + loadOpenAIAutoSwitchSettings, + updateOpenAIAutoSwitchSettings, +} from "../api/openai-auto-switch"; // API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = @@ -68,6 +73,13 @@ const OS_AWARE: Record = { const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); const PROMPT = "Can Unsloth Studio do API calling?"; +// Auto-switch demo: a second call naming a different downloaded GGUF so the +// example shows that the model field selects which model serves. +// A placeholder the user replaces with one of their downloaded GGUFs. A fixed +// repo is usually not one they have, so the resolver would fall through and the +// demo would keep serving the current model instead of switching. +const SWITCH_MODEL = "your-other-downloaded-GGUF"; +const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; // Sampling/thinking knobs for the "+ advanced" examples. @@ -163,13 +175,19 @@ function winBody(model: string, variant: Variant): string { return JSON.stringify(body); } +// A leading comment (valid in both bash and PowerShell) noting the model field +// selects the served model when auto-switch is on. +const SWITCH_NOTE = + '# "Switch model by request" is on: set "model" to any downloaded GGUF to switch.\n'; + function curlUnix( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `curl ${base}/v1/chat/completions \\ + return `${autoSwitch ? SWITCH_NOTE : ""}curl ${base}/v1/chat/completions \\ -H "Authorization: Bearer ${key}" \\ -H "Content-Type: application/json" \\ -d '${shSingle(curlBodyPretty(model, variant))}'`; @@ -181,8 +199,9 @@ function curlWindows( key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { - return `$body = '${psSingle(winBody(model, variant))}' + return `${autoSwitch ? SWITCH_NOTE : ""}$body = '${psSingle(winBody(model, variant))}' Set-Content -Path body.json -Value $body -Encoding ascii curl.exe ${base}/v1/chat/completions \` -H "Authorization: Bearer ${key}" \` @@ -190,11 +209,29 @@ curl.exe ${base}/v1/chat/completions \` -d "@body.json"`; } +// A second OpenAI call naming a different downloaded GGUF: with auto-switch on, +// Studio loads it before serving, so the model field selects the served model. +function pythonSwitchDemo(): string { + return ` + +# "Switch model by request" is on: replace the model below with another GGUF you +# have downloaded and Studio loads it before serving. Unknown names keep serving +# the current model. +response = client.chat.completions.create( + model=${j(SWITCH_MODEL)}, + messages=[{"role": "user", "content": ${j(SWITCH_PROMPT)}}], + stream=True, +) +for chunk in response: + print(chunk.choices[0].delta.content or "", end="")`; +} + function pythonSnippet( base: string, key: string, model: string, variant: Variant, + autoSwitch: boolean, ): string { // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = @@ -241,7 +278,7 @@ response = client.chat.completions.create( messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody} stream=True, ) -${loop}`; +${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; } function buildSnippets( @@ -249,15 +286,16 @@ function buildSnippets( key: string, model: string, os: Os, + autoSwitch: boolean, ): Record { const curl = os === "windows" ? curlWindows : curlUnix; return { - curl: curl(base, key, model, "plain"), - python: pythonSnippet(base, key, model, "plain"), - curlTools: curl(base, key, model, "tools"), - pythonTools: pythonSnippet(base, key, model, "tools"), - curlAdvanced: curl(base, key, model, "advanced"), - pythonAdvanced: pythonSnippet(base, key, model, "advanced"), + curl: curl(base, key, model, "plain", autoSwitch), + python: pythonSnippet(base, key, model, "plain", autoSwitch), + curlTools: curl(base, key, model, "tools", autoSwitch), + pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + curlAdvanced: curl(base, key, model, "advanced", autoSwitch), + pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), }; } @@ -346,12 +384,31 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + // null while loading; the same setting the General tab exposes (shared cache). + const [autoSwitch, setAutoSwitch] = useState( + null, + ); + const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); + useEffect(() => { + let cancelled = false; + void loadOpenAIAutoSwitchSettings() + .then((s) => { + if (!cancelled) setAutoSwitch(s); + }) + .catch(() => { + // Best-effort: leave the toggle off if the setting can't be read. + }); + return () => { + cancelled = true; + }; + }, []); + const model = useLoadedModelName(); // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; @@ -361,9 +418,10 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( - () => buildSnippets(base, key, model, os), - [base, key, model, os], + () => buildSnippets(base, key, model, os, autoSwitchOn), + [base, key, model, os, autoSwitchOn], ); const osAware = OS_AWARE[lang]; @@ -385,6 +443,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { writeUseTunnelPref(next); }; + // Same setting as the General tab; persist optimistically and revert on failure + // so the examples reflect the live model-switch behavior. + const handleToggleAutoSwitch = (next: boolean) => { + const idle = autoSwitch?.autoUnloadIdleSeconds ?? 0; + setAutoSwitch((prev) => (prev ? { ...prev, enabled: next } : prev)); + setSavingAutoSwitch(true); + void updateOpenAIAutoSwitchSettings(next, idle) + .then(setAutoSwitch) + .catch(() => { + setAutoSwitch((prev) => (prev ? { ...prev, enabled: !next } : prev)); + }) + .finally(() => setSavingAutoSwitch(false)); + }; + const handleCopyUrl = async () => { if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) { setCopiedUrl(true); @@ -398,6 +470,41 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.usageExamples")}
+ {/* Same setting as the General tab; surfaced here so the request `model` + actually switches the served model, which the examples below show. */} +
+
+ + + {t("settings.general.modelAutoSwitch.enable")} + + + + + + + {t("settings.general.modelAutoSwitch.enableDescription")} + + +
+
{cloudflareUrl ? (
@@ -412,9 +519,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {/* Only when not launched with --secure: the raw 0.0.0.0 port is still globally reachable, so point the user at --secure. */} - {!secure ? ( + {secure ? null : ( - +
{/* Always rendered (dimmed when off) so toggling never changes the row height and shifts the code block below. */} diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index b3f26c808e..247d5040fb 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -48,6 +48,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -528,6 +529,8 @@ export function GeneralTab() { + + diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index f3c9eef44e..27dda5193a 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -142,6 +142,22 @@ export const en = { loadError: "Failed to load Helper LLM settings.", saveError: "Failed to save Helper LLM settings.", }, + modelAutoSwitch: { + sectionTitle: "Model auto-switch (OpenAI API)", + enable: "Switch model by request", + enableDescription: + "When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.", + idleUnload: "Idle auto-unload", + idleUnloadDescription: + "Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.", + idleNeedsEnable: + "Turn on Switch model by request so an unloaded model reloads on next use.", + idleActiveViaEnv: + "Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.", + loadError: "Failed to load model auto-switch settings.", + saveError: "Failed to save model auto-switch settings.", + idleError: "Enter a whole number of seconds (0 or more).", + }, previewSharing: { sectionTitle: "Preview sharing", enableLabel: "Public preview links",