From 58c2ec1ebdd4564b4ceda8e602e55f001f590808 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 16 Jun 2026 06:17:54 -0700 Subject: [PATCH 01/26] Studio: Xet-primary model downloads with automatic HTTP fallback on stall (#6372) * Studio: add shared Xet-primary download helper with HTTP stall fallback Xet is the fast default transport in huggingface_hub, but a stalled Xet transfer hangs with no progress and no exception, and a blocked native thread cannot be killed. The safetensors inference path already recovers (subprocess watchdog + respawn with HF_HUB_DISABLE_XET=1); the GGUF and training paths do not. Add a reusable helper that the in-process paths can adopt. utils/hf_xet_fallback.py: - DownloadStallError (moved here from core/inference/orchestrator.py, which now imports it; behavior unchanged, still a RuntimeError subclass). - get_hf_download_state / start_watchdog: a no-progress watchdog built on the sparse-aware hub.utils.hf_cache_state helpers; fires only while a .incomplete is present and the on-disk byte total is unchanged for stall_timeout. - hf_hub_download_with_xet_fallback: cached files short-circuit; otherwise the download runs in a spawn child (own process group) supervised by the watchdog. On a stall it kills the child, makes the partial safe for HTTP via prepare_cache_for_transport, and respawns once with HF_HUB_DISABLE_XET=1. Cancel and deterministic errors (auth/missing/disk) propagate without a fallback. Tests cover the watchdog state machine, the transport decision logic, and a regression lock that HF_HUB_DISABLE_XET is honored in a fresh interpreter. * Studio: route GGUF Chat-Mode downloads through the Xet->HTTP fallback The GGUF load path (_download_gguf main+shards, _download_companion_gguf for mmproj/MTP) called a bare blocking hf_hub_download with no recovery, so a Xet stall hung the Chat-Mode load with no fallback. Route those three calls through hf_hub_download_with_xet_fallback: Xet stays primary, HTTP is used only if Xet stalls, per-file so finished shards stay cached. The existing _cancel_event is threaded through, the Cancelled sentinel is preserved, and companions stay best-effort (a terminal stall is swallowed to None). Cached files short-circuit in the helper with no subprocess, so the fast path is unchanged. The two offline mmproj tests are repointed from huggingface_hub.hf_hub_download to the new call boundary (the helper) since the download now goes through it. * Studio: recover a stalled training model-load via Xet->HTTP respawn Training runs in a spawn subprocess and FastModel.from_pretrained downloads internally, so the download cannot be wrapped per-file like GGUF. Instead the worker now watches the HF cache during the model-load phase (emitting model_load_started / model_load_completed and a stall event), and the parent recovers a stall by terminating the worker and respawning it once with HF_HUB_DISABLE_XET=1. worker.py: set HF_HUB_DISABLE_XET=1 before any HF import when the parent passes disable_xet (respawn), and wrap trainer.load_model with start_watchdog. training.py: plumb disable_xet through the config; track the model-load window; on a first-load stall arm a one-shot respawn (handled on the exiting pump thread, so no pump self-join) that preserves the DB run row (history is not duplicated) and re-runs the load over HTTP. A second stall, or a stall outside model-load, surfaces as a normal error. W&B init happens after model-load, so a pre-load respawn cannot duplicate it; the dataset is re-formatted in the new worker. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: add gated test-only fault-injection hook for the Xet stall path UNSLOTH_HF_XET_FORCE_STALL=1 makes the Xet download attempt write a partial blob and hang, so the no-progress watchdog and the HTTP fallback can be exercised end to end against a real repo (never set in production). Used to verify recovery on real models: a forced Xet stall on a 5.37GB Qwen3.5-35B-A3B shard triggered the watchdog and the HTTP retry downloaded the correct file (sha256 verified). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten Xet-fallback comments and consolidate its tests Trim docstrings and inline comments across the Xet->HTTP fallback code to the non-obvious why (spawn-not-thread, killpg-not-getpgid, the sparse-partial HTTP-resume hazard); drop comments that merely restate the code. Verified comment-only with an AST signature check. Merge the three helper-level test files (watchdog, transport policy, and the HF_HUB_DISABLE_XET regression lock) into tests/test_hf_xet_fallback.py, and prefer the real structlog over a bare stub so test collection order cannot leak an incomplete module to others that log at import. Full backend suite: 3455 passed, 14 pre-existing flash-attn failures only. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 34 +- studio/backend/core/inference/orchestrator.py | 8 +- studio/backend/core/training/training.py | 125 ++++++ studio/backend/core/training/worker.py | 48 ++- .../test_gguf_xet_fallback_integration.py | 160 +++++++ studio/backend/tests/test_hf_xet_fallback.py | 352 +++++++++++++++ .../tests/test_offline_gguf_cache_fallback.py | 14 +- .../tests/test_training_xet_fallback.py | 209 +++++++++ studio/backend/utils/hf_xet_fallback.py | 405 ++++++++++++++++++ 9 files changed, 1323 insertions(+), 32 deletions(-) create mode 100644 studio/backend/tests/test_gguf_xet_fallback_integration.py create mode 100644 studio/backend/tests/test_hf_xet_fallback.py create mode 100644 studio/backend/tests/test_training_xet_fallback.py create mode 100644 studio/backend/utils/hf_xet_fallback.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 0f4abe04f5..7b6c0b0668 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -51,6 +51,7 @@ from core.tool_healing import ( strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) @@ -2996,7 +2997,7 @@ class LlamaCppBackend: any time; checks it between each shard download. """ try: - from huggingface_hub import hf_hub_download + import huggingface_hub # noqa: F401 -- presence check only except ImportError: raise RuntimeError( "huggingface_hub is required for HF model loading. " @@ -3140,19 +3141,23 @@ class LlamaCppBackend: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") dl_start = time.monotonic() - local_path = hf_hub_download( - repo_id = hf_repo, - filename = gguf_filename, - token = hf_token, + # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = self._cancel_event, + on_status = lambda m: logger.info(m), ) for shard in gguf_extra_shards: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download( - repo_id = hf_repo, - filename = shard, - token = hf_token, + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = self._cancel_event, ) except RuntimeError as e: if "Cancelled" in str(e): @@ -3213,12 +3218,13 @@ class LlamaCppBackend: return None try: - from huggingface_hub import hf_hub_download logger.info(f"Downloading {label}: {hf_repo}/{target}") - return hf_hub_download( - repo_id = hf_repo, - filename = target, - token = hf_token, + # Same policy; companions are best-effort (caller below swallows failures to None). + return hf_hub_download_with_xet_fallback( + hf_repo, + target, + hf_token, + cancel_event = self._cancel_event, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4fa48a2d88..6b6deb1265 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -30,15 +30,15 @@ from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union from utils.hardware import prepare_gpu_selection +# Re-exported from the shared helper so GGUF, training, and inference share one +# type; kept importable here for backwards compatibility. +from utils.hf_xet_fallback import DownloadStallError + logger = get_logger(__name__) _CTX = mp.get_context("spawn") -class DownloadStallError(RuntimeError): - """Raised when the worker reports no download progress for too long.""" - - # Dispatcher timeout constants (seconds) _DISPATCH_READ_TIMEOUT = 30.0 _DISPATCH_POLL_INTERVAL = 0.5 diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index d69246f584..0c3d42a2fe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -217,6 +217,12 @@ class TrainingBackend: self._db_config: Optional[dict] = None self._db_started_at: Optional[str] = None + # Xet -> HTTP model-load fallback state (config kept for the respawn). + self._last_full_config: Optional[dict] = None + self._in_model_load: bool = False + self._xet_fallback_used: bool = False + self._needs_xet_respawn: bool = False + logger.info("TrainingBackend initialized (subprocess mode)") # ------------------------------------------------------------------ @@ -311,6 +317,8 @@ class TrainingBackend: "trust_remote_code": kwargs.get("trust_remote_code", False), "gpu_ids": kwargs.get("gpu_ids"), "s3_config": kwargs.get("s3_config"), + # Flipped to True only by the HTTP-fallback respawn after a stall. + "disable_xet": kwargs.get("disable_xet", False), } # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. @@ -386,6 +394,11 @@ class TrainingBackend: self._db_total_steps_set = False self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() + # Start each job Xet-first; keep config so a stall can respawn over HTTP. + self._last_full_config = config + self._in_model_load = False + self._xet_fallback_used = False + self._needs_xet_respawn = False # Assign subprocess handles after state reset. self._event_queue = event_queue @@ -446,6 +459,97 @@ class TrainingBackend: output_dir, ) + def _handle_stall_event(self, event: dict) -> None: + """A worker reported a no-progress download stall. + + On the first model-load, terminate the worker so the pump loop respawns it + over HTTP. A later stall (already on HTTP, or outside model-load) surfaces + as an error instead. + """ + msg = event.get("message", "Download stalled") + with self._lock: + recover = self._in_model_load and not self._xet_fallback_used + proc = self._proc + if recover: + self._xet_fallback_used = True + self._needs_xet_respawn = True + self._progress.status_message = ( + "Model download stalled on Xet; retrying over HTTP..." + ) + else: + self._progress.error = self._progress.error or ( + "Model download stalled even over HTTP -- check your network connection" + ) + if recover: + logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg) + else: + logger.error("Training download stalled with no further fallback: %s", msg) + # Terminate either way so the pump loop proceeds (respawn or finalize). + if proc is not None and proc.is_alive(): + proc.terminate() + + def _respawn_worker_disable_xet(self) -> None: + """Respawn the worker once with HF_HUB_DISABLE_XET=1 after a model-load + stall. Runs on the exiting pump thread, reaps the terminated worker, and + starts a fresh worker + pump. DB/progress run-state is preserved so the + history row is not duplicated; the new worker re-formats and loads over HTTP. + """ + config = self._last_full_config + if config is None: + logger.error("Cannot respawn training worker: no stored config") + return + + with self._lock: + old_proc = self._proc + if old_proc is not None: + old_proc.join(timeout = 5.0) + if old_proc.is_alive(): + old_proc.kill() + old_proc.join(timeout = 2.0) + + config = {**config, "disable_xet": True} + self._last_full_config = config + logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall") + + from .worker import run_training_process + + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + new_proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + new_proc.start() + except Exception: + logger.error("Failed to respawn training subprocess", exc_info = True) + with self._lock: + self._progress.is_training = False + self._progress.error = "Failed to recover stalled model download" + self._ensure_db_run_created() + self._finalize_run_in_db( + status = "error", + error_message = "Failed to recover stalled model download", + ) + return + + logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._in_model_load = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = new_proc + self._pump_thread = new_pump + new_pump.start() + def is_training_active(self) -> bool: """Check if training is currently active.""" with self._lock: @@ -566,6 +670,14 @@ class TrainingBackend: for e in self._drain_queue(self._event_queue): self._handle_event(e) + # Model-load stall: respawn over HTTP instead of finalizing as failure. + # Runs on THIS exiting pump thread and starts a fresh pump (never joins + # the current thread); DB run-state is preserved. + if self._needs_xet_respawn: + self._needs_xet_respawn = False + self._respawn_worker_disable_xet() + return + # Mark done if no explicit complete/error was received. with self._lock: if self._progress.is_training: @@ -597,6 +709,19 @@ class TrainingBackend: db_action: Optional[str] = None db_action_kwargs: dict = {} + # Model-load lifecycle + stall recovery (no DB metrics); handled first. + if etype == "model_load_started": + with self._lock: + self._in_model_load = True + return + if etype == "model_load_completed": + with self._lock: + self._in_model_load = False + return + if etype == "stall": + self._handle_stall_event(event) + return + with self._lock: if etype == "progress": self._progress.step = event.get("step", self._progress.step) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index b106aa1298..65533b6946 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1987,6 +1987,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # before imports + # HTTP-fallback respawn: disable Xet before any huggingface_hub import (the + # var is read at import time). Mirrors core/inference/worker.py. + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + print( + "Xet transport disabled for this training worker (HF_HUB_DISABLE_XET=1).", + file = sys.stderr, + flush = True, + ) + # Offline auto-detect: skip ~25s of HF retries per call when DNS is dead. if "HF_HUB_OFFLINE" not in os.environ: import socket as _socket @@ -2706,18 +2719,33 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> cpt_trains_embeddings = False # ── 4c. Load training model (uses VRAM — dataset already formatted) ── + # Watchdog lets the parent recover a stalled Xet download via respawn. _send_status(event_queue, "Loading model...") - success = trainer.load_model( - model_name = model_name, - max_seq_length = config["max_seq_length"], - load_in_4bit = config["load_in_4bit"], - full_finetuning = not use_lora, - hf_token = hf_token, - is_dataset_image = config.get("is_dataset_image", False), - is_dataset_audio = config.get("is_dataset_audio", False), - trust_remote_code = config.get("trust_remote_code", False), - gpu_ids = config.get("resolved_gpu_ids"), + from utils.hf_xet_fallback import start_watchdog + + event_queue.put({"type": "model_load_started", "ts": time.time()}) + _load_watchdog_stop = start_watchdog( + repo_ids = [model_name], + on_stall = lambda msg: event_queue.put( + {"type": "stall", "message": msg, "ts": time.time()} + ), + xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) + try: + success = trainer.load_model( + model_name = model_name, + max_seq_length = config["max_seq_length"], + load_in_4bit = config["load_in_4bit"], + full_finetuning = not use_lora, + hf_token = hf_token, + is_dataset_image = config.get("is_dataset_image", False), + is_dataset_audio = config.get("is_dataset_audio", False), + trust_remote_code = config.get("trust_remote_code", False), + gpu_ids = config.get("resolved_gpu_ids"), + ) + finally: + _load_watchdog_stop.set() + event_queue.put({"type": "model_load_completed", "ts": time.time()}) if not success or trainer.should_stop: if trainer.should_stop: event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) diff --git a/studio/backend/tests/test_gguf_xet_fallback_integration.py b/studio/backend/tests/test_gguf_xet_fallback_integration.py new file mode 100644 index 0000000000..cbcc73847e --- /dev/null +++ b/studio/backend/tests/test_gguf_xet_fallback_integration.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration: GGUF Chat-Mode downloads route through the Xet->HTTP helper, +preserving cancellation and the best-effort companion contract. No GPU, no +network, no real subprocess (the helper is patched). +""" + +from __future__ import annotations + +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Heavy-dep stubbing; prefer the real structlog so a bare stub never leaks to +# later modules that log at import time. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +try: + import structlog # noqa: F401 +except ImportError: + sys.modules["structlog"] = _types.ModuleType("structlog") +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + "HTTPStatusError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Request = type("Request", (), {}) + _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None}) + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **k: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules.setdefault("httpx", _httpx_stub) + +from huggingface_hub import constants as hf_constants + +from core.inference.llama_cpp import LlamaCppBackend +from utils.hf_xet_fallback import DownloadStallError + +REPO = "unsloth/vision-GGUF" + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + return tmp_path + + +def _build_cache( + root: Path, + repo_id: str, + files: dict[str, int], + sha: str = "a" * 40, +) -> Path: + repo_dir = root / f"models--{repo_id.replace('/', '--')}" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + snap = repo_dir / "snapshots" / sha + snap.mkdir(parents = True, exist_ok = True) + for rel, size in files.items(): + (snap / rel).write_bytes(b"\0" * size) + return snap + + +def test_companion_routes_through_helper(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + captured = {} + + def fake_helper( + repo_id, + filename, + token = None, + **kwargs, + ): + captured["filename"] = filename + captured["cancel_event"] = kwargs.get("cancel_event") + return f"/fake/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out == "/fake/mmproj-vision-F16.gguf" + # _cancel_event must be threaded through so /unload can abort the download. + assert captured["cancel_event"] is backend._cancel_event + + +def test_companion_swallows_terminal_stall_to_none(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + + def stalling_helper( + repo_id, + filename, + token = None, + **kwargs, + ): + raise DownloadStallError("both transports stalled") + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", stalling_helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out is None, "a companion download is best-effort; a terminal stall must not raise" + + +def test_companion_cancelled_skips_download(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + backend._cancel_event.set() + called = {"n": 0} + + def helper( + repo_id, + filename, + token = None, + **kwargs, + ): + called["n"] += 1 + return "/should-not-happen" + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out is None + assert called["n"] == 0, "a cancelled load must not start a companion download" diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py new file mode 100644 index 0000000000..39ecebd328 --- /dev/null +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP +transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. +CPU-only, no network, no real subprocess (the per-attempt download seam is +monkeypatched). +""" + +from __future__ import annotations + +import subprocess +import sys +import threading +import time +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub heavy/unavailable deps before importing the module under test. Use the +# real structlog when present; a bare stub left in sys.modules would break later +# modules that log at import time. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +try: + import structlog # noqa: F401 +except ImportError: + sys.modules["structlog"] = _types.ModuleType("structlog") + +import huggingface_hub +from huggingface_hub import constants as hf_constants + +import utils.hf_xet_fallback as xf + + +# --------------------------------------------------------------------------- # +# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. +# --------------------------------------------------------------------------- # +REPO = "ztest/xet-watchdog" + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + return tmp_path + + +def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: + d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" + d.mkdir(parents = True, exist_ok = True) + return d + + +def _wait( + predicate, + timeout: float = 2.0, + step: float = 0.02, +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(step) + return predicate() + + +def test_constant_incomplete_fires_stall(hf_cache): + blobs = _blobs_dir(hf_cache) + (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 + ) + try: + assert _wait( + lambda: len(calls) >= 1, timeout = 3.0 + ), "watchdog never fired on a constant-size .incomplete" + finally: + stop.set() + assert "stalled" in calls[0].lower() + + +def test_growing_incomplete_never_stalls(hf_cache): + blobs = _blobs_dir(hf_cache) + part = blobs / "growing.incomplete" + part.write_bytes(b"\0" * 1024) + + grow_stop = threading.Event() + + def _grow(): + size = 1024 + while not grow_stop.wait(0.05): + size += 4096 + part.write_bytes(b"\0" * size) + + grower = threading.Thread(target = _grow, daemon = True) + grower.start() + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 + ) + try: + time.sleep(1.0) # well past stall_timeout, but bytes keep growing + assert calls == [], "watchdog fired despite continuous progress" + finally: + stop.set() + grow_stop.set() + + +def test_no_incomplete_never_stalls(hf_cache): + blobs = _blobs_dir(hf_cache) + (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 + ) + try: + time.sleep(0.8) + assert calls == [], "watchdog fired with no active .incomplete" + finally: + stop.set() + + +def test_stall_fires_at_most_once(hf_cache): + blobs = _blobs_dir(hf_cache) + (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) + + calls: list[str] = [] + stop = xf.start_watchdog( + repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 + ) + try: + assert _wait(lambda: len(calls) >= 1, timeout = 3.0) + time.sleep(0.6) # keep ticking; must not fire again + assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" + finally: + stop.set() + + +def test_get_state_empty_cache(hf_cache): + assert xf.get_hf_download_state([REPO]) == (0, False) + + +def test_get_state_absent_cache_root(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) + assert xf.get_hf_download_state([REPO]) == (0, False) + + +def test_get_state_skips_local_paths(hf_cache): + # Filesystem paths are not HF repo IDs and must be ignored without error. + assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) + + +def test_get_state_sparse_aware(hf_cache): + blobs = _blobs_dir(hf_cache) + sparse = blobs / "sparse.incomplete" + with open(sparse, "wb") as f: + f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks + st = sparse.stat() + if getattr(st, "st_blocks", 0) == 0: + pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") + total, has_incomplete = xf.get_hf_download_state([REPO]) + assert has_incomplete is True + assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" + + +# --------------------------------------------------------------------------- # +# Transport policy: cached short-circuit, cancel, error propagation, and the +# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. +# --------------------------------------------------------------------------- # +DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" + + +@pytest.fixture(autouse = True) +def _no_real_cache_hit(monkeypatch): + """Default: the cached probe misses; tests override it to force a hit.""" + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + + +class _FakeAttempt: + """Records calls to the download seam and returns scripted results.""" + + def __init__(self, results): + self._results = list(results) + self.calls = [] + + def __call__( + self, + repo_id, + filename, + token, + *, + repo_type, + disable_xet, + cancel_event, + stall_timeout, + interval, + grace_period, + on_status, + ): + self.calls.append( + _types.SimpleNamespace( + repo_id = repo_id, + filename = filename, + disable_xet = disable_xet, + repo_type = repo_type, + ) + ) + return self._results[len(self.calls) - 1] + + +def _install(monkeypatch, results): + fake = _FakeAttempt(results) + monkeypatch.setattr(xf, "_run_download_attempt", fake) + return fake + + +def test_cached_file_short_circuits(monkeypatch, tmp_path): + cached = tmp_path / "cached.gguf" + cached.write_bytes(b"\0" * 8) + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) + fake = _install(monkeypatch, []) # must not be called + + out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert out == str(cached) + assert fake.calls == [], "spawned a download for an already-cached file" + + +def test_cancel_before_start_raises_no_attempt(monkeypatch): + fake = _install(monkeypatch, []) + ev = threading.Event() + ev.set() + with pytest.raises(RuntimeError, match = "Cancelled"): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) + assert fake.calls == [] + + +def test_nonstall_error_propagates_without_fallback(monkeypatch): + fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) + with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" + assert fake.calls[0].disable_xet is False + + +def test_immediate_success_uses_xet_only(monkeypatch): + prepared = [] + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", + lambda *a, **k: prepared.append(a), + ) + fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) + out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert out == "/cache/model.gguf" + assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False + assert prepared == [], "no cache prep should run when Xet succeeds first try" + + +def test_stall_then_http_fallback_succeeds(monkeypatch): + prepared = [] + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", + lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), + ) + fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) + + out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert out == "/cache/model.gguf" + assert len(fake.calls) == 2 + assert fake.calls[0].disable_xet is False # Xet first + assert fake.calls[1].disable_xet is True # HTTP fallback + assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" + + +def test_second_stall_raises_download_stall_error(monkeypatch): + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None + ) + fake = _install(monkeypatch, [("stall", None), ("stall", None)]) + with pytest.raises(xf.DownloadStallError): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert len(fake.calls) == 2 + + +def test_cancelled_midattempt_raises_no_fallback(monkeypatch): + fake = _install(monkeypatch, [("cancelled", None)]) + with pytest.raises(RuntimeError, match = "Cancelled"): + xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) + assert len(fake.calls) == 1 + + +def test_per_file_independent_fallback(monkeypatch): + """A stalled shard falls back; a sibling shard that succeeds does not.""" + monkeypatch.setattr( + "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None + ) + fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) + assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" + assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" + assert [c.disable_xet for c in fake.calls] == [False, False, True] + + +# --------------------------------------------------------------------------- # +# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect +# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). +# --------------------------------------------------------------------------- # +def _safe_path() -> str: + import os + return os.environ.get("PATH", "") + + +def test_disable_xet_constant_set_in_fresh_interpreter(): + code = ( + "from huggingface_hub import constants as c; " + "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" + ) + proc = subprocess.run( + [sys.executable, "-c", code], + env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, + capture_output = True, + text = True, + ) + assert proc.returncode == 0, ( + f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " + f"(rc={proc.returncode}): {proc.stderr}" + ) + + +def test_default_leaves_xet_enabled(): + code = ( + "from huggingface_hub import constants as c; " + "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" + ) + proc = subprocess.run( + [sys.executable, "-c", code], + env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET + capture_output = True, + text = True, + ) + assert proc.returncode == 0, ( + f"without the env var, constants.HF_HUB_DISABLE_XET was not False " + f"(rc={proc.returncode}): {proc.stderr}" + ) diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index a9e0b98247..c00e98ad7e 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -639,17 +639,20 @@ class TestDownloadMmprojOfflineCacheFallback: raise OSError("offline") def fake_download( - *, repo_id, filename, token = None, + **kwargs, ): # Echo back so the test can verify the cache-resolved filename return f"/fake/cache/{repo_id}/{filename}" with ( patch("huggingface_hub.list_repo_files", boom_list), - patch("huggingface_hub.hf_hub_download", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_mmproj( hf_repo = "unsloth/vision-GGUF", @@ -675,17 +678,20 @@ class TestDownloadMmprojOfflineCacheFallback: captured = {} def fake_download( - *, repo_id, filename, token = None, + **kwargs, ): captured["filename"] = filename return f"/fake/{filename}" with ( patch("huggingface_hub.list_repo_files", boom_list), - patch("huggingface_hub.hf_hub_download", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): backend._download_mmproj( hf_repo = "unsloth/vision-GGUF", diff --git a/studio/backend/tests/test_training_xet_fallback.py b/studio/backend/tests/test_training_xet_fallback.py new file mode 100644 index 0000000000..b4a3864334 --- /dev/null +++ b/studio/backend/tests/test_training_xet_fallback.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Parent-side training Xet->HTTP fallback: a model-load stall respawns the +worker once with Xet disabled, preserving the DB run row. Driven via +_handle_event with a fake spawn context; no GPU, no network, no real subprocess. +""" + +from __future__ import annotations + +import contextlib +import logging +import queue +import sys +import threading +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the heavy module-level imports of core/training/training.py so it imports +# under CPU-only/no-network, then restore them (see the restore loop below). +_SAVED: dict = {} + + +def _stub(name, mod): + _SAVED[name] = sys.modules.get(name) + sys.modules[name] = mod + + +_lg = _types.ModuleType("loggers") +_lg.get_logger = lambda name: logging.getLogger(name) +_stub("loggers", _lg) +_stub("structlog", _types.ModuleType("structlog")) +_mpl = _types.ModuleType("matplotlib") +_plt = _types.ModuleType("matplotlib.pyplot") +_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation +_mpl.pyplot = _plt +_stub("matplotlib", _mpl) +_stub("matplotlib.pyplot", _plt) +_hw = _types.ModuleType("utils.hardware") +_hw.prepare_gpu_selection = lambda *a, **k: (None, None) +_stub("utils.hardware", _hw) +_npl = _types.ModuleType("utils.native_path_leases") +_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext() +_npl.run_without_native_path_secret = lambda fn: fn +_stub("utils.native_path_leases", _npl) +_pth = _types.ModuleType("utils.paths") +_pth.outputs_root = lambda *a, **k: "/tmp/outputs" +_stub("utils.paths", _pth) + +import core.training.training as training_mod +from core.training.training import TrainingBackend + +# Restore every stubbed module so this file never pollutes the shared session: a +# leaked bare ``structlog`` (no ``get_logger``) would break every later module +# that logs at import. training_mod already bound the stubs it needs at runtime. +for _name in ( + "loggers", + "structlog", + "matplotlib", + "matplotlib.pyplot", + "utils.hardware", + "utils.native_path_leases", + "utils.paths", +): + _prev = _SAVED.get(_name) + if _prev is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = _prev + + +@pytest.fixture(autouse = True) +def _stub_worker_module(): + """Stub ``core.training.worker`` so the respawn's lazy import of the + torch-heavy worker is never required.""" + prev = sys.modules.get("core.training.worker") + stub = _types.ModuleType("core.training.worker") + stub.run_training_process = lambda **kwargs: None + sys.modules["core.training.worker"] = stub + yield + if prev is None: + sys.modules.pop("core.training.worker", None) + else: + sys.modules["core.training.worker"] = prev + + +class _FakeProc: + def __init__(self, **kwargs): + self._alive = True + self.pid = 4321 + self.kwargs = kwargs + + def start(self): + pass + + def is_alive(self): + return self._alive + + def terminate(self): + self._alive = False + + def kill(self): + self._alive = False + + def join(self, timeout = None): + self._alive = False + + +class _FakeQueue: + def put(self, *a, **k): + pass + + def get(self, *a, **k): + raise queue.Empty + + +class _FakeCtx: + def __init__(self): + self.spawned: list = [] + + def Queue(self): + return _FakeQueue() + + def Process(self, **kwargs): + self.spawned.append(kwargs) + return _FakeProc(**kwargs) + + +def _backend_mid_load(): + b = TrainingBackend() + b._last_full_config = {"model_name": "org/model", "disable_xet": False, "hf_token": "tok"} + b._in_model_load = True + b._xet_fallback_used = False + proc = _FakeProc() + b._proc = proc + return b, proc + + +def test_stall_during_load_arms_respawn_and_terminates_worker(): + b, proc = _backend_mid_load() + b._handle_event({"type": "stall", "message": "no progress for 180s"}) + assert b._needs_xet_respawn is True + assert b._xet_fallback_used is True + assert proc.is_alive() is False, "stalled worker must be terminated" + + +def test_respawn_uses_disable_xet_and_preserves_run_row(monkeypatch): + b, _ = _backend_mid_load() + b._handle_event({"type": "stall", "message": "x"}) + + fake_ctx = _FakeCtx() + monkeypatch.setattr(training_mod, "_CTX", fake_ctx) + monkeypatch.setattr(b, "_pump_loop", lambda: None) # neutralize the new pump + created = {"n": 0} + finalized = {"n": 0} + monkeypatch.setattr( + b, "_ensure_db_run_created", lambda: created.__setitem__("n", created["n"] + 1) + ) + monkeypatch.setattr( + b, "_finalize_run_in_db", lambda **k: finalized.__setitem__("n", finalized["n"] + 1) + ) + + b._respawn_worker_disable_xet() + + assert len(fake_ctx.spawned) == 1, "respawn must start exactly one worker" + cfg = fake_ctx.spawned[0]["kwargs"]["config"] + assert cfg["disable_xet"] is True, "respawned worker must run with Xet disabled" + assert cfg["model_name"] == "org/model" + assert created["n"] == 0, "respawn must not recreate the DB run row" + assert finalized["n"] == 0, "a successful respawn must not finalize the run as error" + + +def test_second_stall_surfaces_error_without_respawn(): + b, proc = _backend_mid_load() + b._xet_fallback_used = True # HTTP fallback already spent + b._handle_event({"type": "stall", "message": "stalled again over http"}) + assert b._needs_xet_respawn is False + assert b._progress.error and "stalled" in b._progress.error.lower() + assert proc.is_alive() is False + + +def test_model_load_completed_disarms_recovery(): + b, _ = _backend_mid_load() + b._handle_event({"type": "model_load_completed"}) + assert b._in_model_load is False + # A stall after the load finished is not a transport stall to recover from. + b._handle_event({"type": "stall", "message": "post-load"}) + assert b._needs_xet_respawn is False + + +def test_model_load_started_arms_recovery_window(): + b = TrainingBackend() + assert b._in_model_load is False + b._handle_event({"type": "model_load_started"}) + assert b._in_model_load is True + + +def test_child_should_disable_xet_truth_table(): + from utils.hf_xet_fallback import child_should_disable_xet + + assert child_should_disable_xet({"disable_xet": True}) is True + assert child_should_disable_xet({"disable_xet": False}) is False + assert child_should_disable_xet({}) is False diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py new file mode 100644 index 0000000000..80f42f92d1 --- /dev/null +++ b/studio/backend/utils/hf_xet_fallback.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. + +Xet (``hf_xet``) is the fast default but can hang with no progress and no +exception, and a blocked native thread cannot be killed. Keep Xet primary; fall +back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` +is read at import time, so the fallback runs in a fresh ``spawn`` child (not a +thread) that sets the env before importing ``huggingface_hub``. Cached files +short-circuit with no child; deterministic errors (401/403/404/disk-full) and +cancellation propagate without a fallback. Mirrors the safetensors inference +recovery in core/inference/{orchestrator,worker}.py. +""" + +from __future__ import annotations + +import multiprocessing as mp +import os +import queue +import signal +import sys +import threading +import time +from typing import Any, Callable, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +_CTX = mp.get_context("spawn") + +# Defaults match the existing inference watchdog and hub shutdown deadline. +DEFAULT_HEARTBEAT_INTERVAL = 30.0 +DEFAULT_STALL_TIMEOUT = 180.0 +DEFAULT_GRACE_PERIOD = 10.0 +_POLL_INTERVAL = 0.5 + + +class DownloadStallError(RuntimeError): + """Raised when no download progress is observed for too long. + + Canonical home; orchestrator.py re-imports it so all paths share one type. + """ + + +def child_should_disable_xet(config: dict) -> bool: + """Single source of truth for the per-worker Xet env flip.""" + return bool(config.get("disable_xet")) + + +def get_hf_download_state( + repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" +) -> Optional[tuple[int, bool]]: + """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. + + Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` + is not mistaken for full-size progress. ``None`` means the state could not be + measured, so callers skip stall logic for that tick. + """ + try: + from hub.utils.hf_cache_state import ( + blob_bytes_present, + has_active_incomplete_blobs, + hf_cache_root, + iter_active_repo_cache_dirs, + ) + + if hf_cache_root() is None: + return (0, False) + + total = 0 + has_incomplete = False + for repo_id in repo_ids or []: + # Skip local paths: HF IDs never start with / . ~ or contain "\". + if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: + continue + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + continue + for blob in blobs_dir.iterdir(): + try: + if blob.is_file(): + total += blob_bytes_present(blob) + except OSError: + pass + if has_active_incomplete_blobs(repo_type, repo_id): + has_incomplete = True + return (total, has_incomplete) + except Exception as e: + logger.debug("Failed to determine HF download state: %s", e) + return None + + +def start_watchdog( + *, + repo_ids: list[str], + on_stall: Callable[[str], None], + repo_type: str = "model", + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + stall_timeout: float = DEFAULT_STALL_TIMEOUT, + xet_disabled: bool = False, + on_heartbeat: Optional[Callable[[str], None]] = None, +) -> threading.Event: + """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a + ``*.incomplete`` is present AND the on-disk size is unchanged for + *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so + post-download init is never misread as a stall. Returns a stop event the + caller sets when the download phase ends. + """ + stop = threading.Event() + transport = "https" if xet_disabled else "xet" + fired = False + + def _beat() -> None: + nonlocal fired + state = get_hf_download_state(repo_ids, repo_type = repo_type) + last_size = state[0] if state is not None else 0 + last_change = time.monotonic() + + while not stop.wait(interval): + state = get_hf_download_state(repo_ids, repo_type = repo_type) + now = time.monotonic() + + if state is None: + if on_heartbeat is not None: + on_heartbeat(f"Downloading ({transport} transport)...") + continue + + current_size, has_incomplete = state + if current_size != last_size: + last_size = current_size + last_change = now + + # Reset unless .incomplete confirms an active download, so model init + # and lock waits are not counted as a stall. + if not has_incomplete: + last_change = now + elif now - last_change >= stall_timeout: + if not fired: + fired = True + on_stall( + f"Download appears stalled ({transport} transport) " + f"-- no progress for {int(now - last_change)}s" + ) + return + + if on_heartbeat is not None: + on_heartbeat(f"Downloading ({transport} transport)...") + + threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() + return stop + + +def _download_child_entry( + *, + repo_id: str, + filename: str, + token: Optional[str], + repo_type: str, + disable_xet: bool, + result_queue: Any, +) -> None: + """Spawn-child entrypoint: download one file and report the result. + + Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, + forms its own process group so the parent can kill the whole transfer, and + never logs the token or signed URLs. + """ + if hasattr(os, "setsid"): + try: + os.setsid() + except OSError: + pass + + if disable_xet: + os.environ["HF_HUB_DISABLE_XET"] = "1" + # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse + # partials a sequential resume cannot safely continue). + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + + # Test-only fault injection (never set in production): stall the Xet attempt + # so the watchdog + HTTP fallback can be exercised against a real repo. + if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": + import time as _t + try: + from huggingface_hub.constants import HF_HUB_CACHE + + blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") + os.makedirs(blobs, exist_ok = True) + with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: + fh.write(b"\0" * 4096) + except OSError: + pass + while True: + _t.sleep(3600) + + try: + from huggingface_hub import hf_hub_download + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + repo_type = repo_type, + token = token, + ) + result_queue.put({"ok": True, "path": path}) + except BaseException as e: # noqa: BLE001 - report every failure to the parent + error = f"{type(e).__name__}: {e}" + try: + from hub.utils.download_registry import scrub_secrets + error = scrub_secrets(error, hf_token = token) + except Exception: + pass + result_queue.put({"ok": False, "error": error}) + + +def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: + """Kill *proc* and its whole process group (Xet may spawn helper procs). + + The child calls ``os.setsid()`` so its pgid equals its pid; signal via + ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a + group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. + """ + pid = proc.pid + + def _signal_group(sig: int) -> None: + if pid is not None and hasattr(os, "killpg"): + try: + os.killpg(pid, sig) + return + except (ProcessLookupError, PermissionError, OSError): + pass + # Windows or pre-setsid: best effort on the single process. + try: + proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() + except Exception: + pass + + _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) + proc.join(timeout = grace_period) + if proc.is_alive(): + _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) + proc.join(timeout = 5.0) + + +def _run_download_attempt( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str, + disable_xet: bool, + cancel_event: Optional[threading.Event], + stall_timeout: float, + interval: float, + grace_period: float, + on_status: Optional[Callable[[str], None]], +) -> tuple[str, Optional[str]]: + """Run one download in a spawn child supervised by the no-progress watchdog. + + Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or + ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. + """ + result_queue: Any = _CTX.Queue() + proc = _CTX.Process( + target = _download_child_entry, + kwargs = dict( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + disable_xet = disable_xet, + result_queue = result_queue, + ), + daemon = True, + ) + proc.start() + + stalled = threading.Event() + stop_watchdog = start_watchdog( + repo_ids = [repo_id], + on_stall = lambda msg: stalled.set(), + repo_type = repo_type, + interval = interval, + stall_timeout = stall_timeout, + xet_disabled = disable_xet, + on_heartbeat = on_status, + ) + + result: Optional[dict] = None + try: + while proc.is_alive(): + if cancel_event is not None and cancel_event.is_set(): + _terminate_process_group(proc, grace_period) + return ("cancelled", None) + if stalled.is_set(): + _terminate_process_group(proc, grace_period) + return ("stall", None) + try: + result = result_queue.get(timeout = _POLL_INTERVAL) + break + except queue.Empty: + continue + else: + # Process exited; drain any result it enqueued. + try: + result = result_queue.get_nowait() + except queue.Empty: + result = None + finally: + stop_watchdog.set() + proc.join(timeout = grace_period) + + if result is None: + return ( + "error", + f"download process for '{repo_id}/{filename}' exited " + f"(code={proc.exitcode}) without a result", + ) + if result.get("ok"): + return ("ok", result["path"]) + return ("error", result.get("error") or "unknown download error") + + +def hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + cancel_event: Optional[threading.Event] = None, + repo_type: str = "model", + stall_timeout: float = DEFAULT_STALL_TIMEOUT, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + grace_period: float = DEFAULT_GRACE_PERIOD, + on_status: Optional[Callable[[str], None]] = None, +) -> str: + """Download a single file with Xet primary and HTTP as a stall-only fallback. + + Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if + *cancel_event* is set, re-raises a deterministic child error unchanged (no + fallback), and raises ``DownloadStallError`` only if BOTH transports stall. + """ + # Finalized blob already cached: return it with no child and no network. + try: + from huggingface_hub import try_to_load_from_cache + cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) + if isinstance(cached, str) and os.path.exists(cached): + return cached + except Exception as e: + logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) + + if cancel_event is not None and cancel_event.is_set(): + raise RuntimeError("Cancelled") + + disable_xet = False + for attempt in range(2): + if disable_xet: + # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume + # over a sparse Xet/hf_transfer partial silently corrupts the blob. + try: + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as e: + logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) + + kind, payload = _run_download_attempt( + repo_id, + filename, + token, + repo_type = repo_type, + disable_xet = disable_xet, + cancel_event = cancel_event, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + ) + + if kind == "ok": + return payload # type: ignore[return-value] + if kind == "cancelled": + raise RuntimeError("Cancelled") + if kind == "error": + # Deterministic failure: the other transport would fail identically. + raise RuntimeError(payload) + # kind == "stall" + if attempt == 0 and not disable_xet: + logger.warning( + "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", + repo_id, + filename, + ) + if on_status is not None: + on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") + disable_xet = True + continue + raise DownloadStallError( + f"Download stalled for '{repo_id}/{filename}' even with " + f"HF_HUB_DISABLE_XET=1 -- check your network connection" + ) + + # Unreachable: the loop either returns or raises on each attempt. + raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") From 14188d6f4592523c62144f20649e07558a37fcb3 Mon Sep 17 00:00:00 2001 From: alkinun Date: Tue, 16 Jun 2026 22:26:13 +0300 Subject: [PATCH 02/26] Add API server monitor in Studio (#5558) * Add Studio API activity monitor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Mark disconnected monitor streams cancelled * Fix monitor lifecycle, parsing, and clock issues for PR #5558 Backend: - Finalize the monitor entry as error on the chat-completions validation reject paths (addresses the codex P2 comment around line 2337). Adds a small _reject helper next to the api_monitor.start so each early-raise in the GGUF tool-passthrough, non-empty-messages, GGUF vision/audio, PIL image decode, and text-only model branches no longer leaves the entry stuck running until eviction. - _monitor_openai_chunk is now defensive about malformed shapes (non-list choices, non-dict choice/delta) so a misbehaving provider chunk does not raise into the streaming generator and abort the user's response. - _monitor_openai_sse_line accepts both data:value and data: value per SSE spec; previously the single-space-only prefix silently dropped tokens from compact emitters. - openai_completions stream now keeps a residual buffer between aiter_bytes chunks so a data: line whose newline lands in the next TCP frame is reconstructed before parsing, instead of being dropped by the per-chunk splitlines call. api_monitor: - duration_ms is derived from time.monotonic anchors and clamped at 0, so NTP / manual clock steps no longer produce negative durations. - finish() and fail() are idempotent: a second call (for example [DONE] arriving after the generator's finally block already ran) no longer moves finished_at or finished_monotonic. - set_usage only derives total_tokens when no authoritative total has been recorded, so a later partial-usage chunk does not clobber a provider-reported total from an earlier chunk. - _trim guards limit < 3 so the slice cannot underflow. Tests: - Adds regression coverage for the new idempotency, total preservation, monotonic clock, and _trim guard behaviour. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Mark monitored stream failures correctly * Fix completions monitor failures * Fix responses stream monitor cleanup * Fix audio input API monitor lifecycle * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix API monitor review follow-ups * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix completions monitor usage accounting * Fix API monitor cancellation and usage gaps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move API monitor into settings * Fix API monitor review followups * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Lazy load API monitor details * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix API monitor review issues * Finalize passthrough monitor on clean EOF * Finalize monitor on chat validation rejects * Fix API monitor review follow-ups * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix API monitor follow-up review issues * Monitor embeddings and tool-call replies * Cover tool-call monitor replies and cancellations * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope API monitor entries by subject * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Finalize monitor entries on cancellation * fix: address PR 5558 CI failures * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: update provider proxy test stub --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: imagineer99 Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/inference/api_monitor.py | 286 ++++ studio/backend/routes/inference.py | 1218 ++++++++++++++--- .../backend/tests/test_anthropic_messages.py | 152 +- studio/backend/tests/test_api_monitor.py | 222 +++ .../tests/test_offline_gguf_cache_fallback.py | 8 +- .../tests/test_openai_tool_passthrough.py | 1206 +++++++++++++++- .../tests/test_responses_tool_passthrough.py | 393 ++++++ .../src/features/chat/api/chat-api.ts | 14 + .../frontend/src/features/chat/types/api.ts | 32 + .../components/api-monitor-console.tsx | 393 ++++++ .../features/settings/tabs/api-keys-tab.tsx | 3 + 11 files changed, 3758 insertions(+), 169 deletions(-) create mode 100644 studio/backend/core/inference/api_monitor.py create mode 100644 studio/backend/tests/test_api_monitor.py create mode 100644 studio/frontend/src/features/settings/components/api-monitor-console.tsx diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py new file mode 100644 index 0000000000..b4ecba1b9d --- /dev/null +++ b/studio/backend/core/inference/api_monitor.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Small in-memory monitor for OpenAI-compatible API traffic.""" + +from __future__ import annotations + +import threading +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Optional + + +_MAX_ENTRIES = 50 +_MAX_PROMPT_CHARS = 12000 +_MAX_REPLY_CHARS = 12000 +_PREVIEW_CHARS = 360 + + +def _trim(text: Optional[str], limit: int) -> str: + if not text: + return "" + if len(text) <= limit: + return text + # Guard against limit < 3 (slice would underflow). + if limit <= 3: + return "..."[:limit] + return text[: limit - 3] + "..." + + +@dataclass +class ApiMonitorEntry: + id: str + endpoint: str + method: str + model: str + prompt: str + status: str + started_at: float + updated_at: float + subject: Optional[str] = None + # Monotonic anchors so duration math survives wall-clock steps (NTP). + started_monotonic: float = 0.0 + finished_monotonic: Optional[float] = None + reply: str = "" + finished_at: Optional[float] = None + context_length: Optional[int] = None + prompt_tokens: Optional[int] = None + completion_tokens: Optional[int] = None + total_tokens: Optional[int] = None + total_tokens_authoritative: bool = False + error: Optional[str] = None + + def snapshot(self, *, include_details: bool = True) -> dict[str, Any]: + duration_ms = None + if self.finished_monotonic is not None: + duration_ms = max( + 0, + int((self.finished_monotonic - self.started_monotonic) * 1000), + ) + elif self.finished_at is not None: + duration_ms = max(0, int((self.finished_at - self.started_at) * 1000)) + context_usage = None + if self.total_tokens is not None and self.context_length: + context_usage = min(1.0, max(0.0, self.total_tokens / self.context_length)) + payload = { + "id": self.id, + "endpoint": self.endpoint, + "method": self.method, + "model": self.model, + "prompt_preview": _trim(self.prompt, _PREVIEW_CHARS), + "reply_preview": _trim(self.reply, _PREVIEW_CHARS), + "prompt_truncated": len(self.prompt) > _PREVIEW_CHARS, + "reply_truncated": len(self.reply) > _PREVIEW_CHARS, + "status": self.status, + "started_at": self.started_at, + "updated_at": self.updated_at, + "finished_at": self.finished_at, + "duration_ms": duration_ms, + "context_length": self.context_length, + "context_usage": context_usage, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.total_tokens, + "error": self.error, + } + if include_details: + payload["prompt"] = self.prompt + payload["reply"] = self.reply + return payload + + +class ApiMonitor: + def __init__(self, max_entries: int = _MAX_ENTRIES): + self._entries: deque[ApiMonitorEntry] = deque() + self._max_entries = max(0, max_entries) + self._lock = threading.Lock() + + def start( + self, + *, + endpoint: str, + method: str, + model: str, + prompt: str, + context_length: Optional[int] = None, + subject: Optional[str] = None, + ) -> str: + now = time.time() + entry = ApiMonitorEntry( + id = f"apireq_{uuid.uuid4().hex[:12]}", + endpoint = endpoint, + method = method, + model = model or "default", + prompt = _trim(prompt, _MAX_PROMPT_CHARS), + status = "running", + started_at = now, + updated_at = now, + subject = subject, + started_monotonic = time.monotonic(), + context_length = context_length, + ) + with self._lock: + self._entries.appendleft(entry) + self._trim_terminal_locked() + return entry.id + + def append_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id or not text: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + entry.reply = _trim(text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_usage( + self, + entry_id: Optional[str], + *, + prompt_tokens: Optional[int] = None, + completion_tokens: Optional[int] = None, + total_tokens: Optional[int] = None, + context_length: Optional[int] = None, + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if prompt_tokens is not None: + entry.prompt_tokens = prompt_tokens + if completion_tokens is not None: + entry.completion_tokens = completion_tokens + if total_tokens is not None: + entry.total_tokens = total_tokens + entry.total_tokens_authoritative = True + elif not entry.total_tokens_authoritative and ( + prompt_tokens is not None or completion_tokens is not None + ): + # Derive only when no authoritative total has been set; + # a later partial chunk must not clobber a provider total. + entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0) + if context_length is not None: + entry.context_length = context_length + entry.updated_at = time.time() + + def finish( + self, + entry_id: Optional[str], + status: str = "completed", + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Idempotent: second call (e.g. [DONE] after the finally block + # already ran) must not move finished_*. + if entry.finished_at is not None: + return + now = time.time() + entry.status = status + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() + + def fail(self, entry_id: Optional[str], error: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if entry.finished_at is not None: + # Already terminal; refresh error text only. + if error: + entry.error = _trim(error, 1000) + return + now = time.time() + entry.status = "error" + entry.error = _trim(error, 1000) + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() + + def snapshot( + self, + *, + include_details: bool = True, + subject: Optional[str] = None, + ) -> list[dict[str, Any]]: + with self._lock: + return [ + entry.snapshot(include_details = include_details) + for entry in self._entries + if subject is None or entry.subject == subject + ] + + def get( + self, + entry_id: str, + *, + subject: Optional[str] = None, + ) -> Optional[dict[str, Any]]: + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return None + if subject is not None and entry.subject != subject: + return None + return entry.snapshot(include_details = True) + + def active_count(self, *, subject: Optional[str] = None) -> int: + with self._lock: + return sum( + 1 + for entry in self._entries + if entry.status == "running" and (subject is None or entry.subject == subject) + ) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + + def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]: + for entry in self._entries: + if entry.id == entry_id: + return entry + return None + + def _trim_terminal_locked(self) -> None: + terminal_seen = 0 + kept: deque[ApiMonitorEntry] = deque() + for entry in self._entries: + if entry.status == "running": + kept.append(entry) + continue + if terminal_seen < self._max_entries: + kept.append(entry) + terminal_seen += 1 + self._entries = kept + + +api_monitor = ApiMonitor() diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d7db0e6f8d..da266bf768 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -875,6 +875,7 @@ from auth.authentication import get_current_subject from state.tool_approvals import resolve_tool_decision from core.inference.key_exchange import decrypt_api_key +from core.inference.api_monitor import api_monitor from core.inference.providers import get_provider_info, get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override @@ -1297,6 +1298,423 @@ def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str logger = get_logger(__name__) +def _monitor_content_text(content) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for part in content: + if isinstance(part, dict): + ptype = part.get("type") + if ptype in ("text", "input_text", "output_text"): + text = part.get("text") + if isinstance(text, str): + parts.append(text) + elif ptype in ("image_url", "input_image", "image"): + parts.append("[image]") + else: + parts.append(f"[{ptype or 'content'}]") + else: + ptype = getattr(part, "type", None) + text = getattr(part, "text", None) + if isinstance(text, str): + parts.append(text) + elif ptype in ("image_url", "input_image", "image"): + parts.append("[image]") + elif ptype: + parts.append(f"[{ptype}]") + return "\n".join(parts) + return str(content) + + +def _monitor_prompt_from_messages(messages) -> str: + lines: list[str] = [] + for msg in messages or []: + role = msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", "") + content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "") + tool_calls = ( + msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None) + ) + text = _monitor_content_text(content) + if tool_calls and not text: + text = "[tool calls]" + if text: + lines.append(f"{role or 'message'}: {text}") + return "\n\n".join(lines) + + +def _monitor_usage( + monitor_id: Optional[str], + usage: Optional[dict], + context_length = None, +): + if not usage: + return + api_monitor.set_usage( + monitor_id, + prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens"), + completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens"), + total_tokens = usage.get("total_tokens"), + context_length = context_length, + ) + + +def _monitor_call_text(name: Any, arguments: Any = None) -> str: + call_name = str(name or "tool") + if arguments is None or arguments == "": + return f"Tool call: {call_name}" + if not isinstance(arguments, str): + args_text = json.dumps(arguments, default = str) + else: + args_text = arguments + if len(args_text) > 500: + args_text = args_text[:497] + "..." + return f"Tool call: {call_name}({args_text})" + + +def _monitor_tool_calls_text(tool_calls: Any) -> str: + if not isinstance(tool_calls, list): + return "" + parts: list[str] = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + continue + fn = tool_call.get("function") or {} + if not isinstance(fn, dict): + fn = {} + name = fn.get("name") or tool_call.get("name") or "tool" + args = fn.get("arguments") + if args is None: + args = tool_call.get("arguments") + parts.append(_monitor_call_text(name, args)) + return "\n".join(parts) + + +def _monitor_openai_chunk( + monitor_id: Optional[str], + data: dict, + context_length = None, +): + if not monitor_id: + return + _monitor_usage(monitor_id, data.get("usage"), context_length) + # Defensive: ignore malformed shapes so the helper never raises into the + # streaming generator and aborts the user's response. + choices = data.get("choices") + if not isinstance(choices, list) or not choices: + return + reply_parts: list[tuple[int, str]] = [] + for idx, choice in enumerate(choices): + if not isinstance(choice, dict): + continue + delta = choice.get("delta") or {} + message = choice.get("message") or {} + content = delta.get("content") if isinstance(delta, dict) else None + if content: + api_monitor.append_reply(monitor_id, content) + continue + if isinstance(delta, dict): + tool_text = _monitor_tool_calls_text(delta.get("tool_calls")) + if tool_text: + api_monitor.append_reply(monitor_id, tool_text) + continue + if isinstance(choice.get("text"), str): + reply_parts.append((idx, choice["text"])) + elif isinstance(message, dict): + text = message.get("content") + if isinstance(text, str): + reply_parts.append((idx, text)) + else: + tool_text = _monitor_tool_calls_text(message.get("tool_calls")) + if tool_text: + reply_parts.append((idx, tool_text)) + if not reply_parts: + return + if len(choices) == 1: + api_monitor.append_reply(monitor_id, reply_parts[0][1]) + return + api_monitor.append_reply( + monitor_id, + "\n\n".join(f"Choice {idx + 1}:\n{text}" for idx, text in reply_parts), + ) + + +def _monitor_openai_error_message(data: dict) -> Optional[str]: + error = data.get("error") + if isinstance(error, dict): + message = error.get("message") + if isinstance(message, str) and message: + return message + return json.dumps(error) + if isinstance(error, str) and error: + return error + return None + + +def _monitor_openai_sse_line( + monitor_id: Optional[str], + raw_line: str, + context_length = None, +) -> Optional[str]: + if not monitor_id: + return None + # SSE spec allows `data:value` and `data: value`; accept both. + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + api_monitor.finish(monitor_id) + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + if isinstance(data, dict): + error_message = _monitor_openai_error_message(data) + if error_message: + api_monitor.fail(monitor_id, error_message) + return "error" + _monitor_openai_chunk(monitor_id, data, context_length) + return None + + +def _monitor_openai_sse_event( + monitor_id: Optional[str], + event: bytes, + context_length = None, +) -> None: + for line in event.decode("utf-8", errors = "ignore").splitlines(): + _monitor_openai_sse_line(monitor_id, line.strip(), context_length) + + +def _monitor_anthropic_usage( + monitor_id: Optional[str], + usage: Optional[dict], + context_length = None, +) -> None: + if not usage: + return + _monitor_usage( + monitor_id, + { + "prompt_tokens": usage.get("input_tokens") or usage.get("prompt_tokens"), + "completion_tokens": usage.get("output_tokens") or usage.get("completion_tokens"), + "total_tokens": usage.get("total_tokens"), + }, + context_length, + ) + + +_ANTHROPIC_MONITOR_TOOL_BLOCKS: dict[str, dict[int, bool]] = {} + + +def _monitor_anthropic_index(data: dict) -> int: + try: + return int(data.get("index") or 0) + except (TypeError, ValueError): + return 0 + + +def _monitor_anthropic_payload( + monitor_id: Optional[str], + data: dict, + context_length = None, +) -> Optional[str]: + if not monitor_id or not isinstance(data, dict): + return None + event_type = data.get("type") + if event_type == "message_start": + message = data.get("message") or {} + if isinstance(message, dict): + _monitor_anthropic_usage(monitor_id, message.get("usage"), context_length) + return None + if event_type == "content_block_start": + content_block = data.get("content_block") or {} + if isinstance(content_block, dict) and content_block.get("type") == "tool_use": + index = _monitor_anthropic_index(data) + _ANTHROPIC_MONITOR_TOOL_BLOCKS.setdefault(monitor_id, {})[index] = False + api_monitor.append_reply(monitor_id, _monitor_call_text(content_block.get("name"))) + return None + if event_type == "content_block_delta": + delta = data.get("delta") or {} + text = delta.get("text") if isinstance(delta, dict) else None + if isinstance(text, str) and text: + api_monitor.append_reply(monitor_id, text) + elif isinstance(delta, dict) and delta.get("type") == "input_json_delta": + index = _monitor_anthropic_index(data) + tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id) or {} + if index in tool_blocks: + if not tool_blocks[index]: + api_monitor.append_reply(monitor_id, "\nInput: ") + tool_blocks[index] = True + partial_json = delta.get("partial_json") + if isinstance(partial_json, str) and partial_json: + api_monitor.append_reply(monitor_id, partial_json) + return None + if event_type == "content_block_stop": + index = _monitor_anthropic_index(data) + tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id) + if tool_blocks is not None: + tool_blocks.pop(index, None) + if not tool_blocks: + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + return None + if event_type == "message_delta": + _monitor_anthropic_usage(monitor_id, data.get("usage"), context_length) + return None + if event_type == "error": + error = data.get("error") or {} + if isinstance(error, dict): + message = error.get("message") or json.dumps(error, default = str) + else: + message = str(error) + api_monitor.fail(monitor_id, message) + return "error" + return None + + +def _monitor_anthropic_sse_line( + monitor_id: Optional[str], + raw_line: str, + context_length = None, +) -> Optional[str]: + if not monitor_id or not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _monitor_anthropic_payload(monitor_id, data, context_length) + + +def _monitor_anthropic_content_blocks(content: Any) -> str: + if not isinstance(content, list): + return "" + parts: list[str] = [] + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(block["text"]) + elif block.get("type") == "tool_use": + parts.append(_monitor_call_text(block.get("name"), block.get("input"))) + return "".join(parts) + + +def _monitor_anthropic_json_response( + response, + monitor_id: Optional[str], + context_length = None, +) -> None: + if not monitor_id: + return + body = getattr(response, "body", b"") + try: + data = json.loads(body.decode("utf-8") if isinstance(body, bytes) else body) + except Exception: + api_monitor.finish(monitor_id) + return + if not isinstance(data, dict): + api_monitor.finish(monitor_id) + return + text = _monitor_anthropic_content_blocks(data.get("content")) + if text: + api_monitor.set_reply(monitor_id, text) + _monitor_anthropic_usage(monitor_id, data.get("usage"), context_length) + api_monitor.finish(monitor_id) + + +def _monitor_anthropic_response( + response, + monitor_id, + context_length = None, + cancel_event = None, +): + if not monitor_id: + return response + body_iterator = getattr(response, "body_iterator", None) + if body_iterator is None: + _monitor_anthropic_json_response(response, monitor_id, context_length) + return response + + async def _monitored_body(): + terminal = False + try: + async for chunk in body_iterator: + text = ( + chunk.decode("utf-8", errors = "ignore") + if isinstance(chunk, (bytes, bytearray)) + else str(chunk) + ) + for line in text.splitlines(): + if ( + _monitor_anthropic_sse_line( + monitor_id, + line.strip(), + context_length, + ) + == "error" + ): + terminal = True + yield chunk + if not terminal: + api_monitor.finish( + monitor_id, + "cancelled" + if cancel_event is not None and cancel_event.is_set() + else "completed", + ) + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + except asyncio.CancelledError: + if cancel_event is not None: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None) + raise + + response.body_iterator = _monitored_body() + return response + + +def _monitor_context_length() -> Optional[int]: + llama_backend = get_llama_cpp_backend() + if getattr(llama_backend, "is_loaded", False): + context_length = _positive_int_or_none(getattr(llama_backend, "context_length", None)) + if context_length is not None: + return context_length + backend = get_inference_backend() + if not backend.active_model_name: + return None + models = getattr(backend, "models", {}) or {} + model_info = models.get(backend.active_model_name, {}) if isinstance(models, dict) else {} + context_length = _positive_int_or_none(model_info.get("context_length")) + if context_length is not None: + return context_length + for candidate in ( + getattr(backend, "context_length", None), + getattr(backend, "max_seq_length", None), + ): + context_length = _positive_int_or_none(candidate) + if context_length is not None: + return context_length + return None + + +def _monitor_active_model() -> Optional[str]: + llama_backend = get_llama_cpp_backend() + if getattr(llama_backend, "is_loaded", False): + return getattr(llama_backend, "model_identifier", None) + backend = get_inference_backend() + return backend.active_model_name + + def _validate_native_gguf_companion( companion_path: str | None, gguf_path: str | None, label: str ) -> None: @@ -2247,6 +2665,35 @@ async def confirm_tool_call( return {"resolved": True} +@studio_router.get("/monitor") +async def get_api_monitor(current_subject: str = Depends(get_current_subject)): + """Return recent OpenAI-compatible API activity for Studio.""" + active_model = _monitor_active_model() + active_requests = api_monitor.active_count(subject = current_subject) + if active_requests: + operating_status = "generating" + elif active_model: + operating_status = "ready" + else: + operating_status = "idle" + return { + "status": operating_status, + "active_model": active_model, + "context_length": _monitor_context_length(), + "active_requests": active_requests, + "entries": api_monitor.snapshot(include_details = False, subject = current_subject), + } + + +@studio_router.get("/monitor/{entry_id}") +async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(get_current_subject)): + """Return full prompt/reply details for one OpenAI-compatible API request.""" + entry = api_monitor.get(entry_id, subject = current_subject) + if entry is None: + raise HTTPException(status_code = 404, detail = "Monitor entry not found") + return entry + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, current_subject: str = Depends(get_current_subject) @@ -3167,7 +3614,9 @@ def _build_external_messages( async def _proxy_to_external_provider( - payload: ChatCompletionRequest, request: Request + payload: ChatCompletionRequest, + request: Request, + current_subject: Optional[str] = None, ) -> StreamingResponse: """ Proxy a chat completion request to an external LLM provider. @@ -3238,6 +3687,16 @@ async def _proxy_to_external_provider( provider_type = provider_type, base_url = base_url, ) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = None, + subject = current_subject, + ) client = ExternalProviderClient( provider_type = provider_type, @@ -3279,14 +3738,29 @@ async def _proxy_to_external_provider( ) try: sent_done = False + stream_failed = False async for line in gen: + monitor_event = _monitor_openai_sse_line(monitor_id, line) + if monitor_event is None: + try: + _monitor_openai_chunk(monitor_id, json.loads(line)) + except Exception: + pass + if monitor_event == "error": + stream_failed = True yield f"{line}\n\n" - if "[DONE]" in line: + if monitor_event == "done": sent_done = True if not sent_done: + if not stream_failed: + api_monitor.finish(monitor_id) yield "data: [DONE]\n\n" + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as exc: logger.error("external_provider.stream_error", error = str(exc)) + api_monitor.fail(monitor_id, _friendly_error(exc)) finally: try: await gen.aclose() @@ -3546,7 +4020,7 @@ async def openai_chat_completions( ) if _wants_multiple_choices(payload): _raise_unsupported_n("external provider chat completions") - return await _proxy_to_external_provider(payload, request) + return await _proxy_to_external_provider(payload, request, current_subject) # Reject a malformed function tool here: it would otherwise reach # llama-server and surface as an opaque 500 "Failed to parse tools". @@ -3595,12 +4069,49 @@ async def openai_chat_completions( # ── Determine which backend is active ───────────────────── # Single-model server: any model name serves the loaded model (drop-in # OpenAI compat), so payload.model is only a fallback label here. + monitor_id = None + + async def _monitored_generate_audio(model_label: str, context_length: Optional[int] = None): + tts_monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + tts_monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model_label, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = context_length, + subject = current_subject, + ) + try: + response = await generate_audio(payload, request) + except asyncio.CancelledError: + api_monitor.finish(tts_monitor_id, "cancelled") + raise + except Exception as e: + api_monitor.fail(tts_monitor_id, _friendly_error(e)) + raise + if isinstance(response, JSONResponse): + try: + body = json.loads(response.body.decode()) + choices = body.get("choices") or [] + message = (choices[0].get("message") or {}) if choices else {} + content = message.get("content") + if isinstance(content, str): + api_monitor.set_reply(tts_monitor_id, content) + except Exception: + pass + api_monitor.finish(tts_monitor_id) + return response + if using_gguf: model_name = llama_backend.model_identifier or payload.model if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") - return await generate_audio(payload, request) + return await _monitored_generate_audio( + model_name, + context_length = llama_backend.context_length, + ) else: backend = get_inference_backend() if not backend.active_model_name: @@ -3616,7 +4127,7 @@ async def openai_chat_completions( # (Whisper is ASR not TTS -- handled below in audio input path) model_info = backend.models.get(backend.active_model_name, {}) if model_info.get("is_audio") and model_info.get("audio_type") != "whisper": - return await generate_audio(payload, request) + return await _monitored_generate_audio(model_name) # ── Whisper without audio: return clear error ── if model_info.get("audio_type") == "whisper" and not payload.audio_base64: @@ -3625,10 +4136,25 @@ async def openai_chat_completions( detail = "Whisper models require audio input. Please upload an audio file.", ) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model_name, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + # ── Audio INPUT path: decode WAV and route to audio input generation ── if payload.audio_base64 and model_info.get("has_audio_input"): - audio_array = _decode_audio_base64(payload.audio_base64) - system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + try: + audio_array = _decode_audio_base64(payload.audio_base64) + system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + except Exception as e: + api_monitor.fail(monitor_id, _friendly_error(e)) + raise cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -3674,16 +4200,20 @@ async def openai_chat_completions( gen = audio_input_generate() _DONE = object() + cancelled = False while True: if cancel_event.is_set(): + cancelled = True break if await request.is_disconnected(): cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") return chunk_text = await asyncio.to_thread(next, gen, _DONE) if chunk_text is _DONE: break if chunk_text: + api_monitor.append_reply(monitor_id, chunk_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -3703,13 +4233,16 @@ async def openai_chat_completions( model = model_name, choices = [ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")], ) + api_monitor.finish(monitor_id, "cancelled" if cancelled else "completed") yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: logger.error(f"Error during audio input streaming: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" finally: _tracker.__exit__(None, None, None) @@ -3724,7 +4257,13 @@ async def openai_chat_completions( }, ) else: - full_text = "".join(audio_input_generate()) + try: + full_text = "".join(audio_input_generate()) + except Exception as e: + api_monitor.fail(monitor_id, _friendly_error(e)) + raise + api_monitor.set_reply(monitor_id, full_text) + api_monitor.finish(monitor_id) response = ChatCompletion( id = completion_id, created = created, @@ -3738,6 +4277,34 @@ async def openai_chat_completions( ) return JSONResponse(content = response.model_dump()) + if monitor_id is None and not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = model_name, + prompt = _monitor_prompt_from_messages(payload.messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + + # Finalize the monitor entry on validation rejection before raising. + def _reject(status_code: int, detail: Any) -> "HTTPException": + if monitor_id is not None: + fail_detail = detail if isinstance(detail, str) else json.dumps(detail, default = str) + api_monitor.fail(monitor_id, fail_detail) + return HTTPException(status_code = status_code, detail = detail) + + def _reject_unsupported_n(path_label: str) -> "HTTPException": + return _reject( + 400, + openai_error_body( + f"n > 1 is not supported for {path_label}.", + status = 400, + code = "unsupported_parameter", + param = "n", + ), + ) + # ── Standard OpenAI function-calling pass-through (GGUF only) ──── # When a client (opencode / Claude Code via OpenAI compat / Cursor / # Continue / ...) sends standard OpenAI `tools` without Studio's @@ -3769,14 +4336,14 @@ async def openai_chat_completions( and (_tools_passthrough or _has_response_format) ): if _wants_multiple_choices(payload): - _raise_unsupported_n("GGUF tool or response_format passthrough") + raise _reject_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: # This path forwards the request verbatim, so the transcoded audio # never gets injected. (The agentic tool loop below does support # audio.) - raise HTTPException( - status_code = 400, - detail = "Audio input is not supported together with guided decoding or client-supplied tools yet.", + raise _reject( + 400, + "Audio input is not supported together with guided decoding or client-supplied tools yet.", ) # Preserve the vision guard from the non-passthrough path below: @@ -3791,9 +4358,9 @@ async def openai_chat_completions( for m in payload.messages ) ): - raise HTTPException( - status_code = 400, - detail = "Image provided but current GGUF model does not support vision.", + raise _reject( + 400, + "Image provided but current GGUF model does not support vision.", ) cancel_event = threading.Event() @@ -3809,21 +4376,20 @@ async def openai_chat_completions( payload, model_name, completion_id, + monitor_id = monitor_id, ) return await _openai_passthrough_non_streaming( llama_backend, payload, model_name, + monitor_id = monitor_id, ) # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages) if not chat_messages: - raise HTTPException( - status_code = 400, - detail = "At least one non-system message is required.", - ) + raise _reject(400, "At least one non-system message is required.") # ── GGUF path: proxy to llama-server /v1/chat/completions ── if using_gguf: @@ -3836,25 +4402,19 @@ async def openai_chat_completions( audio_format = "wav" if payload.audio_base64: if not getattr(llama_backend, "_has_audio_input", False): - raise HTTPException( - status_code = 400, - detail = "Audio provided but current GGUF model does not support audio input.", + raise _reject( + 400, + "Audio provided but current GGUF model does not support audio input.", ) if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: - raise HTTPException( - status_code = 413, - detail = "Audio file is too large (max ~25 MB).", - ) + raise _reject(413, "Audio file is too large (max ~25 MB).") try: audio_b64, audio_format = await asyncio.to_thread( _prepare_audio_for_llama, payload.audio_base64 ) except Exception as e: logger.warning("Audio decode failed: %s", e, exc_info = True) - raise HTTPException( - status_code = 400, - detail = "Could not decode the provided audio file.", - ) + raise _reject(400, "Could not decode the provided audio file.") gguf_messages, _ = _openai_messages_for_gguf_chat( payload, @@ -3918,9 +4478,9 @@ async def openai_chat_completions( # Bypass Permissions suppresses confirm, so the stream requirement # (the gate needs streaming to prompt) no longer applies. if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream: - raise HTTPException( - status_code = 400, - detail = openai_error_body( + raise _reject( + 400, + openai_error_body( "confirm_tool_calls requires stream=true for local tool execution.", status = 400, code = "invalid_request_error", @@ -3928,7 +4488,7 @@ async def openai_chat_completions( ), ) if _wants_multiple_choices(payload): - _raise_unsupported_n("GGUF tool chat completions") + raise _reject_unsupported_n("GGUF tool chat completions") # ── Tool-use system prompt nudge ────────────────────── _nudge = _build_tool_action_nudge( tools = tools_to_use, @@ -4038,6 +4598,7 @@ async def openai_chat_completions( break if await request.is_disconnected(): cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") return event = await asyncio.to_thread(next, gen, _tool_sentinel) @@ -4086,6 +4647,7 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4121,16 +4683,22 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stream_usage, _monitor_context_length()) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: import traceback tb = traceback.format_exc() logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: @@ -4179,7 +4747,7 @@ async def openai_chat_completions( if payload.stream: if _wants_multiple_choices(payload): - _raise_unsupported_n("streaming GGUF chat completions") + raise _reject_unsupported_n("streaming GGUF chat completions") _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() @@ -4212,6 +4780,7 @@ async def openai_chat_completions( break if await request.is_disconnected(): cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") return cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) if cumulative is _gguf_sentinel: @@ -4236,6 +4805,7 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4272,13 +4842,19 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stream_usage, _monitor_context_length()) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: logger.error(f"Error during GGUF streaming: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: @@ -4300,6 +4876,7 @@ async def openai_chat_completions( _n = payload.n or 1 _choices = [] + _monitor_replies = [] _prompt_tokens = 0 _sum_completion = 0 _prompt_details = None @@ -4325,6 +4902,7 @@ async def openai_chat_completions( finish_reason = _clamp_finish_reason(completion_finish), ) ) + _monitor_replies.append(full_text) if completion_usage: # The prompt is shared across all n choices, so count its # tokens ONCE (OpenAI bills only generated tokens for each @@ -4346,10 +4924,27 @@ async def openai_chat_completions( prompt_tokens_details = _prompt_tokens_details(_prompt_details), ), ) + monitor_reply = full_text + if _n > 1: + monitor_reply = "\n\n".join( + f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies) + ) + api_monitor.set_reply(monitor_id, monitor_reply) + _monitor_usage( + monitor_id, + { + "prompt_tokens": _prompt_tokens, + "completion_tokens": _sum_completion, + "total_tokens": _prompt_tokens + _sum_completion, + }, + _monitor_context_length(), + ) + api_monitor.finish(monitor_id) return JSONResponse(content = response.model_dump()) except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) # An over-context prompt makes llama-server return 400; map any # upstream 4xx to a 400 client error rather than leaking a 500. _cls = _classify_llama_generation_error(e) @@ -4469,9 +5064,9 @@ async def openai_chat_completions( # Bypass Permissions suppresses confirm, so the stream requirement # (the gate needs streaming to prompt) no longer applies. if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream: - raise HTTPException( - status_code = 400, - detail = openai_error_body( + raise _reject( + 400, + openai_error_body( "confirm_tool_calls requires stream=true for local tool execution.", status = 400, code = "invalid_request_error", @@ -4593,6 +5188,7 @@ async def openai_chat_completions( if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") return event = await asyncio.to_thread(next, gen, _sf_tool_sentinel) @@ -4627,6 +5223,7 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4667,17 +5264,23 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") raise except Exception: backend.reset_generation_state() # Generic wire message; full trace stays in the log (CWE-209: # transformers/torch errors may leak paths). logger.exception("safetensors tool stream error") + api_monitor.fail(monitor_id, "An internal error occurred.") error_chunk = { "error": { "message": "An internal error occurred.", @@ -4721,6 +5324,11 @@ async def openai_chat_completions( return full_text content_text = await asyncio.to_thread(_drain_to_text) + api_monitor.set_reply(monitor_id, content_text) + _stats = _sf_stats_holder.get("stats") + if _stats: + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") response = ChatCompletion( id = completion_id, created = created, @@ -4733,10 +5341,16 @@ async def openai_chat_completions( ], ) return JSONResponse(content = response.model_dump()) + except asyncio.CancelledError: + cancel_event.set() + backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception: backend.reset_generation_state() # CWE-209: generic detail; full trace in log. logger.exception("safetensors tool completion error") + api_monitor.fail(monitor_id, "An internal error occurred.") raise HTTPException( status_code = 500, detail = "An internal error occurred.", @@ -4830,11 +5444,13 @@ async def openai_chat_completions( if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") return new_text = cumulative[len(prev_text) :] prev_text = cumulative if not new_text: continue + api_monitor.append_reply(monitor_id, new_text) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -4876,15 +5492,21 @@ async def openai_chat_completions( ) if usage_line is not None: yield usage_line + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish( + monitor_id, "cancelled" if cancel_event.is_set() else "completed" + ) yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() backend.reset_generation_state() + api_monitor.finish(monitor_id, "cancelled") raise except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = { "error": { "message": _friendly_error(e), @@ -4923,11 +5545,17 @@ async def openai_chat_completions( ) ], ) + api_monitor.set_reply(monitor_id, full_text) + _stats = stats_holder.get("stats") + if _stats: + _monitor_usage(monitor_id, _stats.get("usage")) + api_monitor.finish(monitor_id) return JSONResponse(content = response.model_dump()) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) + api_monitor.fail(monitor_id, _friendly_error(e)) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) @@ -5137,6 +5765,19 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) + raw_prompt = body.get("prompt", "") + if isinstance(raw_prompt, list): + prompt_text = "\n".join(str(part) for part in raw_prompt) + else: + prompt_text = str(raw_prompt) + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = str(body.get("model") or llama_backend.model_identifier or "default"), + prompt = prompt_text, + context_length = llama_backend.context_length, + subject = current_subject, + ) if is_stream: @@ -5164,10 +5805,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S resp = await _send_stream_with_preheader_cancel(client, req, request = request) if resp is None: + api_monitor.finish(monitor_id, "cancelled") return if resp.status_code != 200: err_bytes = await resp.aread() err_text = err_bytes.decode("utf-8", errors = "replace") + api_monitor.fail(monitor_id, err_text[:500]) raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}") disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, disconnect_event) @@ -5184,26 +5827,49 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge buffer += chunk while b"\n\n" in buffer: event, buffer = buffer.split(b"\n\n", 1) + _monitor_openai_sse_event( + monitor_id, + event, + llama_backend.context_length, + ) out = _cmpl_stream_event_out(event, _include_usage) if out is not None: yield out + b"\n\n" if not disconnect_event.is_set() and buffer: + _monitor_openai_sse_event( + monitor_id, + buffer, + llama_backend.context_length, + ) out = _cmpl_stream_event_out(buffer, _include_usage) if out is not None: # Re-add the SSE separator the split consumed, so a final # event arriving without a trailing blank line is still # terminated for the client's parser. yield out + b"\n\n" + if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") + return + api_monitor.finish(monitor_id) except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not disconnect_event.is_set(): logger.error("openai_completions stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") return + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + disconnect_event.set() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as e: if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return logger.error("openai_completions stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") return @@ -5231,15 +5897,28 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge return StreamingResponse(_stream(), media_type = "text/event-stream") else: - async with httpx.AsyncClient() as client: - resp = await client.post( - target_url, - json = body, - timeout = _llama_non_streaming_generation_timeout(), - ) + try: + async with httpx.AsyncClient() as client: + resp = await client.post( + target_url, + json = body, + timeout = _llama_non_streaming_generation_timeout(), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as e: + api_monitor.fail(monitor_id, _friendly_error(e)) + raise if resp.status_code != 200: + api_monitor.fail(monitor_id, resp.text[:500]) raise _openai_passthrough_error(resp.status_code, resp.text) + try: + _monitor_openai_chunk(monitor_id, resp.json(), llama_backend.context_length) + except Exception: + pass + api_monitor.finish(monitor_id) return Response( content = _rewrite_cmpl_id(resp.content), @@ -5272,9 +5951,43 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get body = await request.json() target_url = f"{llama_backend.base_url}/v1/embeddings" + raw_input = body.get("input", "") + if isinstance(raw_input, list): + prompt_text = "\n".join(str(part) for part in raw_input) + else: + prompt_text = str(raw_input) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = str(body.get("model") or llama_backend.model_identifier or "default"), + prompt = prompt_text, + context_length = llama_backend.context_length, + subject = current_subject, + ) - async with httpx.AsyncClient() as client: - resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S) + try: + async with httpx.AsyncClient() as client: + resp = await client.post( + target_url, + json = body, + timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + if resp.status_code != 200: + api_monitor.fail(monitor_id, resp.text[:500]) + else: + try: + _monitor_usage(monitor_id, resp.json().get("usage"), _monitor_context_length()) + except Exception: + pass + api_monitor.finish(monitor_id) return Response( content = resp.content, status_code = resp.status_code, @@ -5788,84 +6501,128 @@ def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]: async def _responses_non_streaming( - payload: ResponsesRequest, messages: list[ChatMessage], request: Request + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, + current_subject: Optional[str] = None, ) -> JSONResponse: """Handle a non-streaming Responses API call.""" chat_req = _build_chat_request(payload, messages, stream = False) - result = await openai_chat_completions(chat_req, request) - - # openai_chat_completions returns a JSONResponse for non-streaming. - if isinstance(result, JSONResponse): - body = json.loads(result.body.decode()) - elif isinstance(result, Response): - body = json.loads(result.body.decode()) - else: - body = result - - choices = body.get("choices", []) - text = "" - reasoning_text = "" - tool_calls: list[dict] = [] - if choices: - msg = choices[0].get("message", {}) or {} - raw_content = msg.get("content", "") or "" - raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) - llama_backend = get_llama_cpp_backend() - reasoning_text, text = _extract_responses_reasoning( - raw_text, - msg.get("reasoning_content"), - parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), - ) - tool_calls = msg.get("tool_calls") or [] - - usage_data = body.get("usage", {}) - input_tokens = usage_data.get("prompt_tokens", 0) - output_tokens = usage_data.get("completion_tokens", 0) - - resp_id = f"resp_{uuid.uuid4().hex[:12]}" - - # Responses API emits each tool call as its own top-level output item, - # plus an optional assistant text message. Emit the text message only when - # the model produced content, so clients expecting a pure tool-call turn - # (finish_reason="tool_calls") don't see a spurious empty message item. - output_items: list[dict] = [] - if reasoning_text and not text and not tool_calls: - text = reasoning_text - if reasoning_text: - output_items.append(_responses_reasoning_output_item(reasoning_text)) - if text: - msg_id = f"msg_{uuid.uuid4().hex[:12]}" - output_items.append( - ResponsesOutputMessage( - id = msg_id, - status = "completed", - role = "assistant", - content = [ResponsesOutputTextContent(text = text)], - ).model_dump() - ) - output_items.extend(_chat_tool_calls_to_responses_output(tool_calls)) - - response = ResponsesResponse( - id = resp_id, - created_at = int(time.time()), - status = "completed", - model = body.get("model", payload.model), - output = output_items, - usage = ResponsesUsage( - input_tokens = input_tokens, - output_tokens = output_tokens, - total_tokens = input_tokens + output_tokens, - ), - temperature = payload.temperature, - top_p = payload.top_p, - max_output_tokens = payload.max_output_tokens, - instructions = payload.instructions, + request_state = getattr(request, "state", None) + if request_state is None: + request_state = type("_RequestState", (), {})() + try: + setattr(request, "state", request_state) + except Exception: + request_state = None + previous_skip_monitor = ( + bool(getattr(request_state, "skip_api_monitor", False)) + if request_state is not None + else False ) - return JSONResponse(content = response.model_dump()) + monitor_id = None + if not previous_skip_monitor: + monitor_id = api_monitor.start( + endpoint = getattr(getattr(request, "url", None), "path", "/v1/responses"), + method = getattr(request, "method", "POST"), + model = payload.model, + prompt = _monitor_prompt_from_messages(messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + if request_state is not None: + request_state.skip_api_monitor = True + + try: + result = await openai_chat_completions(chat_req, request) + + # openai_chat_completions returns a JSONResponse for non-streaming. + if isinstance(result, JSONResponse): + body = json.loads(result.body.decode()) + elif isinstance(result, Response): + body = json.loads(result.body.decode()) + else: + body = result + + choices = body.get("choices", []) + text = "" + reasoning_text = "" + tool_calls: list[dict] = [] + if choices: + msg = choices[0].get("message", {}) or {} + raw_content = msg.get("content", "") or "" + raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content) + llama_backend = get_llama_cpp_backend() + reasoning_text, text = _extract_responses_reasoning( + raw_text, + msg.get("reasoning_content"), + parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend), + ) + tool_calls = msg.get("tool_calls") or [] + + usage_data = body.get("usage", {}) + input_tokens = usage_data.get("prompt_tokens", 0) + output_tokens = usage_data.get("completion_tokens", 0) + + resp_id = f"resp_{uuid.uuid4().hex[:12]}" + + # Responses API emits each tool call as its own top-level output item, + # plus an optional assistant text message. Emit the text message only when + # the model produced content, so clients expecting a pure tool-call turn + # (finish_reason="tool_calls") don't see a spurious empty message item. + output_items: list[dict] = [] + if reasoning_text and not text and not tool_calls: + text = reasoning_text + if reasoning_text: + output_items.append(_responses_reasoning_output_item(reasoning_text)) + if text: + msg_id = f"msg_{uuid.uuid4().hex[:12]}" + output_items.append( + ResponsesOutputMessage( + id = msg_id, + status = "completed", + role = "assistant", + content = [ResponsesOutputTextContent(text = text)], + ).model_dump() + ) + output_items.extend(_chat_tool_calls_to_responses_output(tool_calls)) + + response = ResponsesResponse( + id = resp_id, + created_at = int(time.time()), + status = "completed", + model = body.get("model", payload.model), + output = output_items, + usage = ResponsesUsage( + input_tokens = input_tokens, + output_tokens = output_tokens, + total_tokens = input_tokens + output_tokens, + ), + temperature = payload.temperature, + top_p = payload.top_p, + max_output_tokens = payload.max_output_tokens, + instructions = payload.instructions, + ) + api_monitor.set_reply(monitor_id, text or _monitor_tool_calls_text(tool_calls)) + _monitor_usage(monitor_id, usage_data, _monitor_context_length()) + api_monitor.finish(monitor_id) + return JSONResponse(content = response.model_dump()) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + finally: + if request_state is not None: + request_state.skip_api_monitor = previous_skip_monitor async def _responses_stream( - payload: ResponsesRequest, messages: list[ChatMessage], request: Request + payload: ResponsesRequest, + messages: list[ChatMessage], + request: Request, + monitor_id: Optional[str] = None, ): """Handle a streaming Responses API call, emitting named SSE events. @@ -6130,9 +6887,11 @@ async def _responses_stream( try: resp = await _send_stream_with_preheader_cancel(client, req, request = request) if resp is None: + api_monitor.finish(monitor_id, "cancelled") return except httpx.RequestError as e: logger.error("responses stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) yield _sse( "response.failed", { @@ -6158,6 +6917,7 @@ async def _responses_stream( resp.status_code, err_text[:500], ) + api_monitor.fail(monitor_id, err_text[:500]) yield _sse( "response.failed", { @@ -6209,6 +6969,11 @@ async def _responses_stream( if usage: input_tokens = usage.get("prompt_tokens", input_tokens) output_tokens = usage.get("completion_tokens", output_tokens) + _monitor_usage( + monitor_id, + usage, + llama_backend.context_length, + ) continue delta = choices[0].get("delta", {}) or {} @@ -6234,6 +6999,7 @@ async def _responses_stream( for event in _ensure_message_open(): yield event full_text += visible_delta + api_monitor.append_reply(monitor_id, visible_delta) yield _sse( "response.output_text.delta", { @@ -6305,9 +7071,18 @@ async def _responses_stream( if usage: input_tokens = usage.get("prompt_tokens", input_tokens) output_tokens = usage.get("completion_tokens", output_tokens) + _monitor_usage( + monitor_id, + usage, + llama_backend.context_length, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not disconnect_event.is_set(): logger.error("responses stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) status_code = 400 if _classify_llama_generation_error(e) is not None else 500 yield _sse( "response.failed", @@ -6316,8 +7091,10 @@ async def _responses_stream( return except Exception as e: if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return logger.error("responses stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) status_code = 400 if _classify_llama_generation_error(e) is not None else 500 yield _sse( "response.failed", @@ -6347,6 +7124,7 @@ async def _responses_stream( pass if disconnect_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return final_reasoning, final_visible = extractor.finish() @@ -6368,6 +7146,7 @@ async def _responses_stream( for event in _ensure_message_open(): yield event full_text += final_visible + api_monitor.append_reply(monitor_id, final_visible) yield _sse( "response.output_text.delta", { @@ -6382,6 +7161,7 @@ async def _responses_stream( for event in _ensure_message_open(): yield event full_text = full_reasoning + api_monitor.set_reply(monitor_id, full_text) yield _sse( "response.output_text.delta", { @@ -6528,6 +7308,7 @@ async def _responses_stream( "arguments": st["arguments"], }, } + api_monitor.append_reply(monitor_id, _monitor_call_text(st["name"], st["arguments"])) yield _sse("response.output_item.done", item_done) # response.completed @@ -6548,6 +7329,7 @@ async def _responses_stream( }, }, } + api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) return StreamingResponse( @@ -6579,8 +7361,28 @@ async def openai_responses( raise HTTPException(status_code = 400, detail = "No input provided.") if payload.stream: - return await _responses_stream(payload, messages, request) - return await _responses_non_streaming(payload, messages, request) + monitor_id = None + if not getattr(request.state, "skip_api_monitor", False): + monitor_id = api_monitor.start( + endpoint = request.url.path, + method = request.method, + model = payload.model, + prompt = _monitor_prompt_from_messages(messages), + context_length = _monitor_context_length(), + subject = current_subject, + ) + try: + return await _responses_stream(payload, messages, request, monitor_id) + except HTTPException as exc: + detail = exc.detail + if not isinstance(detail, str): + detail = json.dumps(detail, default = str) + api_monitor.fail(monitor_id, detail) + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + return await _responses_non_streaming(payload, messages, request, current_subject) # ===================================================================== @@ -6900,14 +7702,67 @@ async def anthropic_messages( and payload.tool_choice.get("disable_parallel_tool_use") ) + monitor_id = None + monitor_context_length = _monitor_context_length() + request_state = getattr(request, "state", None) + if not getattr(request_state, "skip_api_monitor", False): + request_url = getattr(request, "url", None) + monitor_id = api_monitor.start( + endpoint = getattr(request_url, "path", "/v1/messages"), + method = getattr(request, "method", "POST"), + model = model_name, + prompt = _monitor_prompt_from_messages(openai_messages), + context_length = monitor_context_length, + subject = current_subject, + ) + + async def _monitored_anthropic(coro): + try: + response = await coro + except asyncio.CancelledError: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + raise + except Exception as exc: + api_monitor.fail(monitor_id, _friendly_error(exc)) + raise + return _monitor_anthropic_response( + response, + monitor_id, + monitor_context_length, + cancel_event, + ) + # ── Client-side pass-through path ───────────────────────── if client_tools: openai_tools = openai_client_tools if payload.stream: - return await _anthropic_passthrough_stream( - request, - cancel_event, + return await _monitored_anthropic( + _anthropic_passthrough_stream( + request, + cancel_event, + llama_backend, + openai_messages, + openai_tools, + temperature, + top_p, + top_k, + payload.max_tokens, + message_id, + model_name, + stop = stop, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + tool_choice = openai_tool_choice, + session_id = payload.session_id, + cancel_id = payload.cancel_id, + disable_parallel_tool_use = _disable_parallel, + ) + ) + return await _monitored_anthropic( + _anthropic_passthrough_non_streaming( llama_backend, openai_messages, openai_tools, @@ -6922,26 +7777,8 @@ async def anthropic_messages( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = openai_tool_choice, - session_id = payload.session_id, - cancel_id = payload.cancel_id, disable_parallel_tool_use = _disable_parallel, ) - return await _anthropic_passthrough_non_streaming( - llama_backend, - openai_messages, - openai_tools, - temperature, - top_p, - top_k, - payload.max_tokens, - message_id, - model_name, - stop = stop, - min_p = min_p, - repetition_penalty = repetition_penalty, - presence_penalty = presence_penalty, - tool_choice = openai_tool_choice, - disable_parallel_tool_use = _disable_parallel, ) if server_tools: @@ -6949,6 +7786,10 @@ async def anthropic_messages( if bool(getattr(payload, "confirm_tool_calls", False)) and not bool( getattr(payload, "bypass_permissions", False) ): + api_monitor.fail( + monitor_id, + "confirm_tool_calls is not supported for Anthropic Messages server tools.", + ) raise HTTPException( status_code = 400, detail = anthropic_error_body( @@ -7009,22 +7850,26 @@ async def anthropic_messages( ) if payload.stream: - return await _anthropic_tool_stream( - request, - cancel_event, + return await _monitored_anthropic( + _anthropic_tool_stream( + request, + cancel_event, + _run_tool_gen, + message_id, + model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, + openai_tools = openai_tools, + disable_parallel_tool_use = _disable_parallel, + ) + ) + return await _monitored_anthropic( + _anthropic_tool_non_streaming( _run_tool_gen, message_id, model_name, - llama_backend = llama_backend, - openai_messages = openai_messages, - openai_tools = openai_tools, disable_parallel_tool_use = _disable_parallel, ) - return await _anthropic_tool_non_streaming( - _run_tool_gen, - message_id, - model_name, - disable_parallel_tool_use = _disable_parallel, ) # ── No-tool path ────────────────────────────────────────── @@ -7043,19 +7888,23 @@ async def anthropic_messages( ) if payload.stream: - return await _anthropic_plain_stream( - request, - cancel_event, + return await _monitored_anthropic( + _anthropic_plain_stream( + request, + cancel_event, + _run_plain_gen, + message_id, + model_name, + llama_backend = llama_backend, + openai_messages = openai_messages, + ) + ) + return await _monitored_anthropic( + _anthropic_plain_non_streaming( _run_plain_gen, message_id, model_name, - llama_backend = llama_backend, - openai_messages = openai_messages, ) - return await _anthropic_plain_non_streaming( - _run_plain_gen, - message_id, - model_name, ) @@ -8106,7 +8955,13 @@ def _build_openai_passthrough_body( async def _openai_passthrough_stream( - request, cancel_event, llama_backend, payload, model_name, completion_id + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, ): """Streaming client-side pass-through for /v1/chat/completions. @@ -8150,6 +9005,7 @@ async def _openai_passthrough_stream( except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) if resp is not None: try: await resp.aclose() @@ -8164,6 +9020,7 @@ async def _openai_passthrough_stream( detail = _friendly_error(e), ) if resp is None: + api_monitor.finish(monitor_id, "cancelled") try: await client.aclose() except Exception: @@ -8205,6 +9062,7 @@ async def _openai_passthrough_stream( await client.aclose() except Exception: pass + api_monitor.fail(monitor_id, err_text[:500]) raise _openai_passthrough_error(upstream_status, err_text) async def _stream(): @@ -8218,6 +9076,7 @@ async def _openai_passthrough_stream( disconnect_watcher = asyncio.create_task( _await_disconnect_then_close(request, resp, cancel_event) ) + monitor_done = False try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -8237,21 +9096,39 @@ async def _openai_passthrough_stream( # relayed byte-for-byte. if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: raw_line = _cap_parallel_tool_calls_sse_line(raw_line) + monitor_event = _monitor_openai_sse_line( + monitor_id, + raw_line, + llama_backend.context_length, + ) # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" - if raw_line[6:].strip() == "[DONE]": + if monitor_event == "done" or raw_line[6:].strip() == "[DONE]": + monitor_done = True break + if not monitor_done: + api_monitor.finish( + monitor_id, + "cancelled" if cancel_event.is_set() else "completed", + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. if not cancel_event.is_set(): + api_monitor.fail(monitor_id, "Stream interrupted") raise + api_monitor.finish(monitor_id, "cancelled") except Exception as e: if cancel_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") return # 200 headers already flushed; errors must go in the SSE body. logger.error("openai passthrough stream error: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: @@ -8294,7 +9171,12 @@ async def _openai_passthrough_stream( raise -async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): +async def _openai_passthrough_non_streaming( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, +): """Non-streaming client-side pass-through for /v1/chat/completions. Returns llama-server's JSON response verbatim so the client sees the native @@ -8317,11 +9199,15 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): json = body, timeout = _llama_non_streaming_generation_timeout(), ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. Surface the # same friendly message the sync chat path emits so operators don't see # a bare 500 with no diagnostic. logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) raise HTTPException( status_code = 502, detail = _friendly_error(e), @@ -8337,6 +9223,7 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): ): _truncate_budget -= 1 continue + api_monitor.fail(monitor_id, resp.text[:500]) raise _openai_passthrough_error(resp.status_code, resp.text) # The guided-decoding fence wraps each choice's JSON content in a @@ -8357,6 +9244,7 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): "openai passthrough non-streaming: response not JSON, relaying raw: %s", exc, ) + api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") changed = False @@ -8394,5 +9282,9 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name): # Nothing mutated: relay the upstream bytes verbatim, skipping a redundant # parse + re-serialize round-trip. if not changed: + _monitor_openai_chunk(monitor_id, data, llama_backend.context_length) + api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") + _monitor_openai_chunk(monitor_id, data, llama_backend.context_length) + api_monitor.finish(monitor_id) return JSONResponse(content = data) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 0ed8de5532..92a87ce045 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -33,6 +33,7 @@ from core.inference.anthropic_compat import ( AnthropicStreamEmitter, AnthropicPassthroughEmitter, ) +from core.inference.api_monitor import ApiMonitor from routes.inference import ( _build_tool_action_nudge, _normalize_anthropic_openai_images, @@ -40,6 +41,7 @@ from routes.inference import ( _anthropic_requested_studio_tools, _anthropic_passthrough_stream, _anthropic_tool_non_streaming, + _monitor_anthropic_sse_line, anthropic_messages, ) from state.tool_policy import reset_tool_policy, set_tool_policy @@ -50,6 +52,46 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/messages", + method = "POST", + model = "m", + prompt = "hi", + ) + + for payload in ( + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_1", + "name": "lookup", + "input": {}, + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "input_json_delta", + "partial_json": '{"query":"weather"}', + }, + }, + {"type": "content_block_stop", "index": 0}, + ): + _monitor_anthropic_sse_line(monitor_id, f"data: {json.dumps(payload)}") + + entry = monitor.get(monitor_id) + assert entry is not None + assert entry["reply"] == 'Tool call: lookup\nInput: {"query":"weather"}' + + # ===================================================================== # Tool nudge tests # ===================================================================== @@ -1385,7 +1427,7 @@ def _mock_backend(monkeypatch, **overrides): def _gen_plain(**kwargs): calls.append(("plain", kwargs)) - yield {"type": "content", "text": "ok"} + yield "ok" def _gen_tools(**kwargs): calls.append(("tools", kwargs)) @@ -1396,6 +1438,8 @@ def _mock_backend(monkeypatch, **overrides): is_vision = False, supports_tools = True, model_identifier = "test-model", + context_length = 4096, + count_chat_tokens = lambda *args, **kwargs: 2, generate_chat_completion = _gen_plain, generate_chat_completion_with_tools = _gen_tools, calls = calls, @@ -1426,6 +1470,112 @@ def _reset_policy(): class TestAnthropicMessagesToolRouting: + class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + @staticmethod + def _consume_response(response): + async def _consume(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + return chunks + + return _drive(_consume()) + + def test_plain_non_streaming_records_api_monitor_entry(self, monkeypatch): + import routes.inference as inf_mod + + _mock_backend(monkeypatch, context_length = 2048) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + payload = _basic_payload() + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + assert response.status_code == 200 + [entry] = monitor.snapshot() + assert entry["endpoint"] == "/v1/messages" + assert entry["status"] == "completed" + assert entry["model"] == "test-model" + assert entry["prompt_preview"] == "user: hi" + assert entry["reply_preview"] == "ok" + assert entry["context_length"] == 2048 + assert monitor.active_count() == 0 + + def test_tool_use_non_streaming_records_api_monitor_reply(self, monkeypatch): + import routes.inference as inf_mod + + def _gen_tools(**_kwargs): + yield { + "type": "tool_start", + "tool_call_id": "call_1", + "tool_name": "lookup", + "arguments": {"query": "weather"}, + } + + _mock_backend( + monkeypatch, + context_length = 2048, + generate_chat_completion_with_tools = _gen_tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + payload = _basic_payload( + enable_tools = True, + tools = [{"type": "web_search_20250305", "name": "web_search"}], + ) + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + assert response.status_code == 200 + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply_preview"] == 'Tool call: lookup({"query": "weather"})' + + def test_plain_streaming_records_active_and_completed_monitor_entry(self, monkeypatch): + import routes.inference as inf_mod + + _mock_backend(monkeypatch, context_length = 2048) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + payload = _basic_payload(stream = True) + + response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + assert monitor.active_count() == 1 + self._consume_response(response) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply_preview"] == "ok" + assert entry["prompt_tokens"] == 2 + assert entry["context_length"] == 2048 + assert monitor.active_count() == 0 + + def test_plain_streaming_pre_response_cancel_finalizes_monitor(self, monkeypatch): + import routes.inference as inf_mod + + async def _cancelled_before_response(*_args, **_kwargs): + raise asyncio.CancelledError() + + _mock_backend(monkeypatch, context_length = 2048) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_anthropic_plain_stream", _cancelled_before_response) + payload = _basic_payload(stream = True) + + with pytest.raises(asyncio.CancelledError): + _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t")) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch): _mock_backend(monkeypatch) payload = _basic_payload( diff --git a/studio/backend/tests/test_api_monitor.py b/studio/backend/tests/test_api_monitor.py new file mode 100644 index 0000000000..f0943cd76c --- /dev/null +++ b/studio/backend/tests/test_api_monitor.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from core.inference.api_monitor import ApiMonitor, _trim + + +def test_api_monitor_tracks_reply_usage_and_context(): + monitor = ApiMonitor(max_entries = 3) + + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "local-model", + prompt = "user: hello", + context_length = 100, + ) + monitor.append_reply(entry_id, "hi") + monitor.append_reply(entry_id, " there") + monitor.set_usage( + entry_id, + prompt_tokens = 4, + completion_tokens = 6, + ) + monitor.finish(entry_id) + + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hi there" + assert entry["total_tokens"] == 10 + assert entry["context_usage"] == 0.1 + assert entry["duration_ms"] is not None + + +def test_api_monitor_summary_omits_full_prompt_and_reply(): + monitor = ApiMonitor(max_entries = 3) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "local-model", + prompt = "p" * 500, + ) + monitor.set_reply(entry_id, "r" * 500) + + [summary] = monitor.snapshot(include_details = False) + assert "prompt" not in summary + assert "reply" not in summary + assert summary["prompt_preview"].endswith("...") + assert summary["reply_preview"].endswith("...") + assert summary["prompt_truncated"] is True + assert summary["reply_truncated"] is True + + detail = monitor.get(entry_id) + assert detail is not None + assert detail["prompt"] == "p" * 500 + assert detail["reply"] == "r" * 500 + + +def test_api_monitor_filters_entries_by_subject(): + monitor = ApiMonitor(max_entries = 3) + alice = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "alice prompt", + subject = "alice", + ) + bob = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "bob prompt", + subject = "bob", + ) + monitor.finish(bob) + + alice_entries = monitor.snapshot(subject = "alice") + assert [entry["id"] for entry in alice_entries] == [alice] + assert monitor.get(bob, subject = "alice") is None + assert monitor.get(bob, subject = "bob")["id"] == bob + assert monitor.active_count(subject = "alice") == 1 + assert monitor.active_count(subject = "bob") == 0 + + +def test_api_monitor_keeps_bounded_recent_history(): + monitor = ApiMonitor(max_entries = 2) + + first = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "first", + ) + second = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "second", + ) + third = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "third", + ) + monitor.finish(first) + monitor.finish(second) + monitor.finish(third) + + entries = monitor.snapshot() + ids = [entry["id"] for entry in entries] + assert ids[0] == third + assert [entry["prompt"] for entry in entries] == ["third", "second"] + assert first not in ids + assert monitor.active_count() == 0 + + +def test_api_monitor_keeps_running_entries_beyond_history_limit(): + monitor = ApiMonitor(max_entries = 1) + + running = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "running", + ) + for prompt in ("done-1", "done-2", "done-3"): + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = prompt, + ) + monitor.finish(entry_id) + + entries = monitor.snapshot() + ids = [entry["id"] for entry in entries] + assert running in ids + assert monitor.active_count() == 1 + + monitor.finish(running) + [entry] = monitor.snapshot() + assert entry["id"] == running + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + +def test_api_monitor_finish_is_idempotent(): + monitor = ApiMonitor(max_entries = 2) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "hi", + ) + monitor.finish(entry_id) + first = monitor.snapshot()[0] + monitor.finish(entry_id) + second = monitor.snapshot()[0] + assert first["finished_at"] == second["finished_at"] + assert first["duration_ms"] == second["duration_ms"] + + +def test_api_monitor_preserves_authoritative_total_tokens(): + monitor = ApiMonitor(max_entries = 2) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "hi", + ) + monitor.set_usage( + entry_id, + prompt_tokens = 10, + completion_tokens = 20, + total_tokens = 33, + ) + # A later partial chunk omitting `total_tokens` must not clobber 33. + monitor.set_usage(entry_id, prompt_tokens = 11) + assert monitor.snapshot()[0]["total_tokens"] == 33 + + +def test_api_monitor_recomputes_derived_total_tokens(): + monitor = ApiMonitor(max_entries = 2) + entry_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "m", + prompt = "hi", + ) + monitor.set_usage(entry_id, prompt_tokens = 10) + assert monitor.snapshot()[0]["total_tokens"] == 10 + + monitor.set_usage(entry_id, completion_tokens = 20) + entry = monitor.snapshot()[0] + assert entry["prompt_tokens"] == 10 + assert entry["completion_tokens"] == 20 + assert entry["total_tokens"] == 30 + + +def test_api_monitor_duration_non_negative_under_clock_step(monkeypatch): + import core.inference.api_monitor as m + + fake_now = [1000.0] + monkeypatch.setattr(m.time, "time", lambda: fake_now[0]) + monitor = ApiMonitor(max_entries = 1) + entry_id = monitor.start( + endpoint = "/x", + method = "POST", + model = "m", + prompt = "hi", + ) + fake_now[0] = 500.0 + monitor.finish(entry_id) + assert monitor.snapshot()[0]["duration_ms"] >= 0 + + +def test_api_monitor_trim_guards_tiny_limit(): + assert _trim("abcdefgh", 2) == ".." + assert _trim("abcdefgh", 0) == "" + assert _trim("abcdefgh", 3) == "..." + assert _trim("abcdefgh", 4) == "a..." + assert _trim("abcdefgh", 100) == "abcdefgh" diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index c00e98ad7e..a2a505f479 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -209,10 +209,10 @@ class TestGgufVariantFileResolution: return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None] def fake_download( - *, repo_id, filename, token = None, + **_kwargs, ): downloaded.append(filename) return f"/fake/{repo_id}/{filename}" @@ -229,7 +229,7 @@ class TestGgufVariantFileResolution: ), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("huggingface_hub.hf_hub_download", fake_download), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), ): out = backend._download_gguf( hf_repo = "ggml-org/models", @@ -256,10 +256,10 @@ class TestGgufVariantFileResolution: return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None] def fake_download( - *, repo_id, filename, token = None, + **_kwargs, ): downloaded.append(filename) return f"/fake/{repo_id}/{filename}" @@ -269,7 +269,7 @@ class TestGgufVariantFileResolution: patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("huggingface_hub.hf_hub_download", fake_download), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), ): out = backend._download_gguf( hf_repo = "org/repo", diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index b2b22f8934..8442d05cdc 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -7,6 +7,7 @@ import os import sys import asyncio import json +import threading from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -27,20 +28,29 @@ from models.inference import ( from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) +from core.inference.api_monitor import ApiMonitor from routes.inference import ( _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, _clamp_finish_reason, + _cmpl_stream_event_out, _coalesce_consecutive_user_turns, _drop_empty_assistant_sentinels, _effective_max_tokens, _extract_content_parts, _friendly_error, _merge_user_content, + _monitor_openai_chunk, + _monitor_openai_sse_event, _openai_messages_for_gguf_chat, + _openai_passthrough_non_streaming, + _openai_passthrough_stream, _openai_stream_usage_chunk, + _proxy_to_external_provider, _set_or_prepend_system_message, + openai_completions, + openai_embeddings, openai_chat_completions, ) from state.tool_policy import reset_tool_policy @@ -325,7 +335,8 @@ class TestChatCompletionRequestToolFields: captured = {} - async def _fake_proxy(payload, request): + async def _fake_proxy(payload, request, current_subject): + assert current_subject == "test-user" captured["stream"] = payload.stream return JSONResponse({"choices": [], "object": "chat.completion"}) @@ -460,6 +471,7 @@ class TestChatCompletionRequestToolFields: supports_tools = False is_vision = False _is_audio = False + context_length = 4096 client = self._v1_client(monkeypatch, _GGUFBackend()) resp = client.post( @@ -473,13 +485,18 @@ class TestChatCompletionRequestToolFields: self._assert_unsupported_n(resp) def test_n_rejected_for_gguf_tools_passthrough_path(self, monkeypatch): + import routes.inference as inference_route + class _GGUFBackend: is_loaded = True model_identifier = "test-gguf" supports_tools = True is_vision = False _is_audio = False + context_length = 4096 + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) client = self._v1_client(monkeypatch, _GGUFBackend()) resp = client.post( "/v1/chat/completions", @@ -498,6 +515,10 @@ class TestChatCompletionRequestToolFields: }, ) self._assert_unsupported_n(resp) + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "n > 1 is not supported" in entry["error"] + assert monitor.active_count() == 0 def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: @@ -540,6 +561,8 @@ class TestChatCompletionRequestToolFields: "_detect_safetensors_features", lambda backend, chat_template: {"supports_tools": True}, ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend()) resp = client.post( "/v1/chat/completions", @@ -556,6 +579,10 @@ class TestChatCompletionRequestToolFields: body = resp.json() assert body["error"]["param"] == "confirm_tool_calls" assert "requires stream=true" in body["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "confirm_tool_calls requires stream=true" in entry["error"] + assert monitor.active_count() == 0 def test_multiturn_tool_loop_messages(self): req = ChatCompletionRequest( @@ -887,6 +914,35 @@ class TestOpenAICompatibilityHelpers: assert usage["completion_tokens"] == 7 assert usage["total_tokens"] == 7 + def test_completion_stream_monitor_reads_usage_before_client_strip(self, monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/completions", + method = "POST", + model = "m", + prompt = "hi", + context_length = 100, + ) + event = ( + b'data: {"id":"chatcmpl-test","choices":[{"text":"done","finish_reason":"stop"}],' + b'"usage":{"prompt_tokens":4,"completion_tokens":6,"total_tokens":10}}\n' + ) + + _monitor_openai_sse_event(monitor_id, event, context_length = 100) + out = _cmpl_stream_event_out(event, include_usage = False) + + assert out is not None + assert b'"usage"' not in out + [entry] = monitor.snapshot() + assert entry["reply"] == "done" + assert entry["prompt_tokens"] == 4 + assert entry["completion_tokens"] == 6 + assert entry["total_tokens"] == 10 + assert entry["context_usage"] == 0.1 + def test_developer_message_preserves_existing_system_prompt(self): payload = ChatCompletionRequest( messages = [ @@ -1168,6 +1224,10 @@ class TestGgufVisionMessages: class TestGgufVisionToolRouting: class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + async def is_disconnected(self): return False @@ -1203,9 +1263,12 @@ class TestGgufVisionToolRouting: is_vision = True, supports_tools = True, model_identifier = "gemma-4-12b-it-GGUF", + context_length = 4096, generate_chat_completion = _plain, generate_chat_completion_with_tools = _tools, ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) payload = ChatCompletionRequest( @@ -1258,9 +1321,12 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = True, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _plain, generate_chat_completion_with_tools = _tools, ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) payload = ChatCompletionRequest( @@ -1292,10 +1358,13 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = True, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _plain, generate_chat_completion_with_tools = _tools, ) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) payload = ChatCompletionRequest( model = "default", @@ -1316,6 +1385,62 @@ class TestGgufVisionToolRouting: ) assert exc.value.status_code == 400 assert "requires stream=true" in exc.value.detail["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "confirm_tool_calls requires stream=true" in entry["error"] + assert monitor.active_count() == 0 + + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): + import routes.inference as inf_mod + + calls = {"count": 0} + + def _generate(**_kwargs): + calls["count"] += 1 + text = f"reply {calls['count']}" + yield text + yield { + "type": "metadata", + "usage": { + "prompt_tokens": 3, + "completion_tokens": calls["count"], + "total_tokens": 3 + calls["count"], + }, + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + n = 2, + messages = [{"role": "user", "content": "two please"}], + ) + + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + body = json.loads(response.body) + + assert [c["message"]["content"] for c in body["choices"]] == ["reply 1", "reply 2"] + [entry] = monitor.snapshot() + assert entry["reply"] == "Choice 1:\nreply 1\n\nChoice 2:\nreply 2" + assert entry["completion_tokens"] == 3 + assert monitor.active_count() == 0 def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1336,6 +1461,7 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = False, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _generate, ) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) @@ -1369,6 +1495,8 @@ class TestGgufVisionToolRouting: import routes.inference as inf_mod seen_seeds = [] + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) def _generate(**kwargs): seen_seeds.append(kwargs.get("seed")) @@ -1388,6 +1516,7 @@ class TestGgufVisionToolRouting: is_vision = False, supports_tools = False, model_identifier = "test-gguf", + context_length = 4096, generate_chat_completion = _generate, ) monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) @@ -1406,6 +1535,1081 @@ class TestGgufVisionToolRouting: assert seen_seeds == expected assert [choice["index"] for choice in body["choices"]] == [0, 1, 2] + assert body["usage"]["prompt_tokens"] == 5 + assert body["usage"]["completion_tokens"] == 21 + [entry] = monitor.snapshot() + assert entry["prompt_tokens"] == 5 + assert entry["completion_tokens"] == 21 + assert entry["total_tokens"] == 26 + + +class TestApiMonitorProviderAndCompletionStreams: + class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + async def is_disconnected(self): + return False + + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyExternalClient: + def __init__(self, **_kwargs): + pass + + async def stream_chat_completion(self, **kwargs): + assert kwargs["stream"] is False + yield json.dumps( + { + "choices": [{"message": {"content": "provider [DONE] reply"}}], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + }, + } + ) + + async def close(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "ExternalProviderClient", DummyExternalClient) + payload = ChatCompletionRequest( + model = "default", + external_model = "gpt-test", + provider_type = "openai", + provider_base_url = "https://api.openai.com/v1", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + response = await _proxy_to_external_provider(payload, self._Request()) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert chunks[-1] == "data: [DONE]\n\n" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "provider [DONE] reply" + assert entry["prompt_tokens"] == 3 + assert entry["completion_tokens"] == 4 + assert entry["total_tokens"] == 7 + + asyncio.run(_run()) + + def test_external_stream_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyExternalClient: + def __init__(self, **_kwargs): + pass + + async def stream_chat_completion(self, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + await asyncio.sleep(3600) + + async def close(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "ExternalProviderClient", DummyExternalClient) + payload = ChatCompletionRequest( + model = "default", + external_model = "gpt-test", + provider_type = "openai", + provider_base_url = "https://api.openai.com/v1", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await _proxy_to_external_provider(payload, self._Request()) + iterator = response.body_iterator + first = await anext(iterator) + assert "hello" in first + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_preheader_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": True} + + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return None + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + response = await openai_completions(Request(), current_subject = "test") + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert chunks == [] + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_stream_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": True} + + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield b'data: {"choices":[{"text":"hello"}]}\n\n' + await asyncio.sleep(3600) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + + response = await openai_completions(Request(), current_subject = "test") + iterator = response.body_iterator + first = await anext(iterator) + assert b"hello" in first + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_non_streaming_post_error_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False} + + class FailingAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **_kwargs): + raise httpx.ConnectError("llama down") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod.httpx, + "AsyncClient", + lambda *args, **kwargs: FailingAsyncClient(), + ) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + with pytest.raises(httpx.ConnectError): + await openai_completions(Request(), current_subject = "test") + + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "Lost connection to the model server" in entry["error"] + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + _monitor_openai_chunk( + monitor_id, + { + "choices": [ + {"text": "first"}, + {"text": "second"}, + ], + "usage": { + "prompt_tokens": 2, + "completion_tokens": 5, + "total_tokens": 7, + }, + }, + 4096, + ) + + entry = monitor.get(monitor_id) + assert entry["reply"] == "Choice 1:\nfirst\n\nChoice 2:\nsecond" + assert entry["prompt_tokens"] == 2 + assert entry["completion_tokens"] == 5 + assert entry["context_length"] == 4096 + + def test_monitor_openai_chunk_records_tool_call_reply(self, monkeypatch): + import routes.inference as inf_mod + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + _monitor_openai_chunk( + monitor_id, + { + "choices": [ + { + "message": { + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"query":"weather"}', + }, + } + ] + } + } + ] + }, + 4096, + ) + + entry = monitor.get(monitor_id) + assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' + + def test_embeddings_request_is_counted_active_and_completed(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/embeddings") + method = "POST" + + async def json(self): + return {"input": ["alpha", "beta"], "model": "embed"} + + class FakeAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **_kwargs): + assert monitor.active_count() == 1 + return httpx.Response( + 200, + json = { + "data": [{"embedding": [0.1]}], + "usage": {"prompt_tokens": 4, "total_tokens": 4}, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod.httpx, + "AsyncClient", + lambda *args, **kwargs: FakeAsyncClient(), + ) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + response = await openai_embeddings(Request(), current_subject = "test") + + assert response.status_code == 200 + [entry] = monitor.snapshot() + assert entry["endpoint"] == "/v1/embeddings" + assert entry["status"] == "completed" + assert entry["prompt_preview"] == "alpha\nbeta" + assert entry["prompt_tokens"] == 4 + assert entry["total_tokens"] == 4 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + await asyncio.sleep(3600) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + first = await anext(iterator) + assert "hello" in first + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class CancellingAsyncClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **_kwargs): + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod.httpx, + "AsyncClient", + lambda *args, **kwargs: CancellingAsyncClient(), + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'] + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + +class TestApiMonitorSafetensorsUsage: + class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + def test_non_streaming_safetensors_records_usage(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyBackend: + active_model_name = "safe-model" + models = {"safe-model": {"context_length": 2048}} + + def generate_chat_response(self, *, stats_holder, **_kwargs): + stats_holder["stats"] = { + "usage": { + "prompt_tokens": 8, + "completion_tokens": 5, + "total_tokens": 13, + } + } + yield "safe reply" + + def reset_generation_state(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, + "_detect_safetensors_features", + lambda *_args, **_kwargs: {"supports_tools": False}, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + body = json.loads(response.body) + + assert body["choices"][0]["message"]["content"] == "safe reply" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "safe reply" + assert entry["prompt_tokens"] == 8 + assert entry["completion_tokens"] == 5 + assert entry["total_tokens"] == 13 + assert entry["context_length"] == 2048 + + asyncio.run(_run()) + + def test_non_streaming_safetensors_tool_cancel_records_cancelled(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + reset_tool_policy() + + class DummyBackend: + active_model_name = "safe-model" + models = {"safe-model": {"context_length": 2048}} + + def generate_chat_response(self, **_kwargs): + raise AssertionError("plain safetensors path should not be used") + + def generate_chat_completion_with_tools( + self, *, cancel_event, stats_holder, **_kwargs + ): + stats_holder["stats"] = { + "usage": { + "prompt_tokens": 8, + "completion_tokens": 5, + "total_tokens": 13, + } + } + yield {"type": "content", "text": "partial"} + cancel_event.set() + yield {"type": "content", "text": "ignored"} + + def reset_generation_state(self): + pass + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, + "_detect_safetensors_features", + lambda *_args, **_kwargs: {"supports_tools": True}, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + enable_tools = True, + enabled_tools = ["web_search"], + cancel_id = "safe-cancel", + ) + + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + body = json.loads(response.body) + + assert body["choices"][0]["message"]["content"] == "partial" + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "partial" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_non_streaming_safetensors_tool_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + reset_tool_policy() + reset_called = False + + class DummyBackend: + active_model_name = "safe-model" + models = {"safe-model": {"context_length": 2048}} + + def generate_chat_response(self, **_kwargs): + raise AssertionError("plain safetensors path should not be used") + + def generate_chat_completion_with_tools(self, **_kwargs): + yield {"type": "content", "text": "unused"} + + def reset_generation_state(self): + nonlocal reset_called + reset_called = True + + async def fake_to_thread(*_args, **_kwargs): + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod.asyncio, "to_thread", fake_to_thread) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend()) + monkeypatch.setattr( + inf_mod, + "_detect_safetensors_features", + lambda *_args, **_kwargs: {"supports_tools": True}, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + enable_tools = True, + enabled_tools = ["web_search"], + cancel_id = "safe-cancel", + ) + + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert reset_called is True + + asyncio.run(_run()) + + +class TestApiMonitorAudioInput: + def _patch_audio_backend(self, monkeypatch, chunks): + import routes.inference as inf_mod + + class DummyAudioBackend: + active_model_name = "audio-model" + models = { + "audio-model": { + "has_audio_input": True, + "audio_type": "audio-input", + } + } + + def generate_audio_input_response(self, **_kwargs): + yield from chunks + + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False), + ) + monkeypatch.setattr( + inf_mod, + "get_inference_backend", + lambda: DummyAudioBackend(), + ) + monkeypatch.setattr( + inf_mod, + "_decode_audio_base64", + lambda _payload: object(), + ) + return inf_mod + + def test_audio_input_non_streaming_records_active_monitor(self, monkeypatch): + async def _run(): + inf_mod = self._patch_audio_backend(monkeypatch, ["hello", " world"]) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "describe this audio")], + audio_base64 = "ZmFrZQ==", + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + response = await openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + body = json.loads(response.body) + + assert body["choices"][0]["message"]["content"] == "hello world" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello world" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_audio_input_streaming_records_monitor_reply(self, monkeypatch): + async def _run(): + inf_mod = self._patch_audio_backend(monkeypatch, ["hello", " world"]) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + async def is_disconnected(): + return False + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "describe this audio")], + audio_base64 = "ZmFrZQ==", + stream = True, + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + is_disconnected = is_disconnected, + ) + + response = await openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + + assert chunks[-1] == "data: [DONE]\n\n" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "hello world" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_non_gguf_tts_auto_route_records_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyTtsBackend: + active_model_name = "tts-model" + models = { + "tts-model": { + "is_audio": True, + "audio_type": "snac", + } + } + + async def fake_generate_audio( + _payload, + _request, + current_subject = None, + ): + return inf_mod.JSONResponse( + content = { + "choices": [ + { + "message": { + "content": "[Generated audio]", + } + } + ] + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend()) + monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "say hello")], + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + response = await inf_mod.openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == ( + "[Generated audio]" + ) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["model"] == "tts-model" + assert entry["reply"] == "[Generated audio]" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_non_gguf_tts_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class DummyTtsBackend: + active_model_name = "tts-model" + models = { + "tts-model": { + "is_audio": True, + "audio_type": "snac", + } + } + + async def fake_generate_audio( + _payload, + _request, + current_subject = None, + ): + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False), + ) + monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend()) + monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "say hello")], + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + with pytest.raises(asyncio.CancelledError): + await inf_mod.openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["model"] == "tts-model" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tts_auto_route_records_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_generate_audio( + _payload, + _request, + current_subject = None, + ): + return inf_mod.JSONResponse( + content = { + "choices": [ + { + "message": { + "content": "[Generated audio]", + } + } + ] + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + _is_audio = True, + model_identifier = "gguf-tts", + context_length = 2048, + ), + ) + monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "say hello")], + ) + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/chat/completions"), + method = "POST", + ) + + await inf_mod.openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["model"] == "gguf-tts" + assert entry["context_length"] == 2048 + assert entry["reply"] == "[Generated audio]" + assert monitor.active_count() == 0 + + asyncio.run(_run()) # ===================================================================== diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 89155c2daf..0bea355668 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -40,6 +40,7 @@ from fastapi import HTTPException from fastapi.responses import JSONResponse from pydantic import ValidationError +from core.inference.api_monitor import ApiMonitor from models.inference import ( ChatMessage, ResponsesFunctionCallInputItem, @@ -781,6 +782,129 @@ class TestResponsesNonStreamingAdapter: assert "" not in body["output"][1]["content"][0]["text"] assert "" not in body["output"][1]["content"][0]["text"] + def test_monitor_records_translated_visible_text(self, monkeypatch): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + assert request.state.skip_api_monitor is True + return JSONResponse( + content = { + "model": "test-model", + "choices": [{"message": {"content": "plananswer"}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/responses"), + method = "POST", + ) + + async def run(): + response = await _responses_non_streaming(payload, messages, request) + return json.loads(response.body.decode()) + + body = asyncio.run(run()) + + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][1]["content"][0]["text"] == "answer" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "answer" + assert entry["prompt_tokens"] == 2 + assert entry["completion_tokens"] == 3 + assert request.state.skip_api_monitor is False + + def test_monitor_records_tool_only_reply(self, monkeypatch): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + assert request.state.skip_api_monitor is True + return JSONResponse( + content = { + "model": "test-model", + "choices": [ + { + "message": { + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"query":"weather"}', + }, + } + ], + } + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 3}, + } + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + payload = ResponsesRequest( + input = "hi", + tools = [{"type": "function", "name": "lookup"}], + ) + messages = [ChatMessage(role = "user", content = "hi")] + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/responses"), + method = "POST", + ) + + async def run(): + response = await _responses_non_streaming(payload, messages, request) + return json.loads(response.body.decode()) + + body = asyncio.run(run()) + + assert body["output"][0]["type"] == "function_call" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' + assert request.state.skip_api_monitor is False + + def test_cancelled_chat_completion_finalizes_monitor(self, monkeypatch): + import routes.inference as inf_mod + + async def fake_chat_completions(chat_req, request): + assert request.state.skip_api_monitor is True + raise asyncio.CancelledError() + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions) + payload = ResponsesRequest(input = "hi") + messages = [ChatMessage(role = "user", content = "hi")] + request = SimpleNamespace( + state = SimpleNamespace(), + url = SimpleNamespace(path = "/v1/responses"), + method = "POST", + ) + + async def run(): + with pytest.raises(asyncio.CancelledError): + await _responses_non_streaming(payload, messages, request) + + asyncio.run(run()) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert request.state.skip_api_monitor is False + def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch): body = self._run_with_message(monkeypatch, {"content": "show x tags"}) @@ -939,6 +1063,275 @@ class TestResponsesStreamAdapter: assert completed["response"]["output"][0]["content"][0]["text"] == "plan" assert completed["response"]["output"][1]["content"][0]["text"] == "33" + def test_usage_only_chunk_updates_monitor(self, monkeypatch): + import routes.inference as inf_mod + + chunks = [ + {"choices": [{"delta": {"content": "33"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + asyncio.run(run()) + + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "33" + assert entry["prompt_tokens"] == 2 + assert entry["completion_tokens"] == 3 + assert entry["total_tokens"] == 5 + assert entry["context_length"] == 4096 + + def test_function_call_chunk_updates_monitor_reply(self, monkeypatch): + import routes.inference as inf_mod + + chunks = [ + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "function": { + "name": "lookup", + "arguments": '{"query":"weather"}', + }, + } + ] + } + } + ] + } + ] + self._install_stream_mock(monkeypatch, chunks) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert self._payloads(lines, "response.output_item.done")[-1]["item"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == 'Tool call: lookup({"query":"weather"})' + + def test_preheader_cancel_finalizes_monitor(self, monkeypatch): + import routes.inference as inf_mod + + self._install_stream_mock(monkeypatch, []) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + + async def fake_send(*_args, **_kwargs): + return None + + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + asyncio.run(run()) + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + def test_stream_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + await asyncio.sleep(3600) + + self._install_stream_mock(monkeypatch, []) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + iterator = response.body_iterator + first = "" + for _ in range(8): + first = await anext(iterator) + if "hello" in first: + break + else: + pytest.fail("stream did not emit text delta") + + pending = asyncio.create_task(anext(iterator)) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_final_visible_text_updates_monitor(self, monkeypatch): + import routes.inference as inf_mod + + class FakeExtractor: + def __init__(self, **_kwargs): + pass + + def feed( + self, + _content, + _reasoning_content = None, + ): + return "", "" + + def finish(self): + return "", "tail" + + self._install_stream_mock(monkeypatch, [{"choices": [{"delta": {"content": ""}}]}]) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_ResponsesReasoningExtractor", FakeExtractor) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "m", + prompt = "hi", + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, + messages, + self._Request(), + monitor_id = monitor_id, + ) + return await self._collect(response) + + lines = asyncio.run(run()) + + assert self._payloads(lines, "response.output_text.delta")[-1]["delta"] == "plan" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plan" + def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch): chunks = [ {"choices": [{"delta": {"content": "show { return parseJsonOrThrow(response); } +export async function getApiMonitor(): Promise { + const response = await authFetch("/api/inference/monitor"); + return parseJsonOrThrow(response); +} + +export async function getApiMonitorEntry(id: string): Promise { + const response = await authFetch( + `/api/inference/monitor/${encodeURIComponent(id)}`, + ); + return parseJsonOrThrow(response); +} + export async function loadModel( payload: LoadModelRequest, ): Promise { diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 2f0e30fdcf..781172fb4f 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -193,6 +193,38 @@ export interface InferenceStatusResponse { spec_fallback_reason?: string | null; } +export interface ApiMonitorEntry { + id: string; + endpoint: string; + method: string; + model: string; + prompt?: string; + reply?: string; + prompt_preview: string; + reply_preview: string; + prompt_truncated: boolean; + reply_truncated: boolean; + status: "running" | "completed" | "cancelled" | "error"; + started_at: number; + updated_at: number; + finished_at?: number | null; + duration_ms?: number | null; + context_length?: number | null; + context_usage?: number | null; + prompt_tokens?: number | null; + completion_tokens?: number | null; + total_tokens?: number | null; + error?: string | null; +} + +export interface ApiMonitorResponse { + status: "idle" | "ready" | "generating"; + active_model?: string | null; + context_length?: number | null; + active_requests: number; + entries: ApiMonitorEntry[]; +} + export interface AudioGenerationResponse { id: string; object: string; diff --git a/studio/frontend/src/features/settings/components/api-monitor-console.tsx b/studio/frontend/src/features/settings/components/api-monitor-console.tsx new file mode 100644 index 0000000000..36d09115ae --- /dev/null +++ b/studio/frontend/src/features/settings/components/api-monitor-console.tsx @@ -0,0 +1,393 @@ +// 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 { cn } from "@/lib/utils"; +import { + ActivityIcon, + ChevronDownIcon, + CircleIcon, + RefreshCwIcon, +} from "lucide-react"; +import { + type ReactElement, + useCallback, + useEffect, + useMemo, + useRef, + useState, +} from "react"; +import { getApiMonitor, getApiMonitorEntry } from "../../chat/api/chat-api"; +import type { ApiMonitorEntry, ApiMonitorResponse } from "../../chat/types/api"; + +const API_INFERENCE_PREFIX_RE = /^\/api\/inference/; +const V1_PREFIX_RE = /^\/v1\//; + +function formatTime(value: number): string { + return new Date(value * 1000).toLocaleTimeString([], { + hour: "2-digit", + minute: "2-digit", + second: "2-digit", + }); +} + +function formatDuration(value?: number | null): string { + if (value == null) { + return "Running"; + } + if (value < 1000) { + return `${value} ms`; + } + return `${(value / 1000).toFixed(value < 10000 ? 1 : 0)} s`; +} + +function formatTokens(entry: ApiMonitorEntry): string { + if (entry.total_tokens != null) { + return `${entry.total_tokens.toLocaleString()} tokens`; + } + if (entry.prompt_tokens != null || entry.completion_tokens != null) { + const prompt = entry.prompt_tokens ?? 0; + const completion = entry.completion_tokens ?? 0; + return `${(prompt + completion).toLocaleString()} tokens`; + } + return "Tokens pending"; +} + +function compactEndpoint(endpoint: string): string { + return endpoint + .replace(API_INFERENCE_PREFIX_RE, "/api") + .replace(V1_PREFIX_RE, "/"); +} + +function statusTone(status: ApiMonitorEntry["status"]): string { + if (status === "running") { + return "text-emerald-500"; + } + if (status === "error") { + return "text-destructive"; + } + if (status === "cancelled") { + return "text-amber-500"; + } + return "text-muted-foreground"; +} + +function UsageBar({ value }: { value?: number | null }): ReactElement | null { + if (value == null) { + return null; + } + const pct = Math.max(0, Math.min(100, Math.round(value * 100))); + return ( +
+
+
+ ); +} + +function MonitorEntry({ + entry, + detail, + expanded, + loading, + onToggle, +}: { + entry: ApiMonitorEntry; + detail?: ApiMonitorEntry; + expanded: boolean; + loading: boolean; + onToggle: () => void; +}): ReactElement { + const hasCurrentDetail = + detail && + detail.status === entry.status && + detail.updated_at >= entry.updated_at; + const prompt = detail?.prompt ?? entry.prompt_preview; + const replyText = hasCurrentDetail + ? detail.error ?? detail.reply ?? entry.error ?? entry.reply_preview + : entry.error ?? entry.reply_preview; + const reply = replyText || (entry.status === "running" ? "Waiting..." : "No reply"); + + return ( +
+ + + {expanded ? ( +
+
+
+
+ Prompt + {entry.prompt_truncated && !detail ? Preview : null} +
+ 
+                {loading && !detail ? "Loading..." : prompt || "No prompt text"}
+
+
+
+
+ Reply + {entry.reply_truncated && !detail ? Preview : null} +
+ 
+                {loading && !detail ? "Loading..." : reply}
+
+
+
+ +
+ {formatTokens(entry)} + {entry.context_length ? ( + <> / {entry.context_length.toLocaleString()} context + ) : null} + +
+
+ ) : null} +
+ ); +} + +export function ApiMonitorConsole(): ReactElement { + const [data, setData] = useState(null); + const [error, setError] = useState(null); + const [refreshing, setRefreshing] = useState(false); + const [expandedIds, setExpandedIds] = useState>(() => new Set()); + const [details, setDetails] = useState>({}); + const [loadingDetails, setLoadingDetails] = useState>( + () => new Set(), + ); + const loadingDetailsRef = useRef>(new Set()); + const detailsRef = useRef>({}); + + const loadMonitor = useCallback(async (): Promise => { + setRefreshing(true); + try { + setData(await getApiMonitor()); + setError(null); + } catch (err: unknown) { + setError(err instanceof Error ? err.message : "Monitor unavailable"); + } finally { + setRefreshing(false); + } + }, []); + + useEffect(() => { + let cancelled = false; + let timer: number | undefined; + + function schedule(): void { + timer = window.setTimeout(poll, 1500); + } + + function poll(): void { + getApiMonitor() + .then((next) => { + if (cancelled) { + return; + } + setData(next); + setError(null); + }) + .catch((err: unknown) => { + if (cancelled) { + return; + } + setError(err instanceof Error ? err.message : "Monitor unavailable"); + }) + .finally(() => { + if (!cancelled) { + schedule(); + } + }); + } + + poll(); + return () => { + cancelled = true; + if (timer !== undefined) { + window.clearTimeout(timer); + } + }; + }, []); + + const statusLabel = data?.status ?? "idle"; + const hasActive = (data?.active_requests ?? 0) > 0; + const entries = useMemo(() => data?.entries ?? [], [data]); + const loadDetail = useCallback( + (id: string): void => { + if (loadingDetailsRef.current.has(id)) { + return; + } + loadingDetailsRef.current.add(id); + setLoadingDetails((prev) => new Set(prev).add(id)); + getApiMonitorEntry(id) + .then((entry) => { + setDetails((prev) => { + const next = { ...prev, [id]: entry }; + detailsRef.current = next; + return next; + }); + }) + .catch(() => { + setDetails((prev) => { + const next = { ...prev }; + delete next[id]; + detailsRef.current = next; + return next; + }); + }) + .finally(() => { + loadingDetailsRef.current.delete(id); + setLoadingDetails((prev) => { + const next = new Set(prev); + next.delete(id); + return next; + }); + }); + }, + [], + ); + + const toggleEntry = useCallback( + (entry: ApiMonitorEntry): void => { + setExpandedIds((prev) => { + const next = new Set(prev); + if (next.has(entry.id)) { + next.delete(entry.id); + } else { + next.add(entry.id); + loadDetail(entry.id); + } + return next; + }); + }, + [loadDetail], + ); + + useEffect(() => { + for (const entry of entries) { + if (!expandedIds.has(entry.id)) { + continue; + } + const cached = detailsRef.current[entry.id]; + if (!cached || cached.status !== entry.status || entry.status === "running") { + loadDetail(entry.id); + } + } + }, [entries, expandedIds, loadDetail]); + + return ( +
+
+
+
+ + {hasActive ? ( + + ) : null} +
+
+

+ API monitor +

+

+ {data?.active_model ?? "No model loaded"} +

+
+
+
+
+ {statusLabel} +
+ +
+
+ +
+ + {(data?.active_requests ?? 0).toLocaleString()} active /{" "} + {entries.length.toLocaleString()} recent + + {data?.context_length ? ( + {data.context_length.toLocaleString()} context + ) : null} +
+ +
+ {error ? ( +
+ {error} +
+ ) : entries.length === 0 ? ( +
+ No API traffic yet +
+ ) : ( +
+ {entries.map((entry) => ( + toggleEntry(entry)} + /> + ))} +
+ )} +
+
+ ); +} diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index 64f3ef6621..e83ecebfc8 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -14,6 +14,7 @@ import { translate, useT } from "@/i18n"; import { AnimatePresence, motion, useReducedMotion } from "motion/react"; import { useCallback, useEffect, useState } from "react"; import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys"; +import { ApiMonitorConsole } from "../components/api-monitor-console"; import { ApiKeyRow } from "../components/api-key-row"; import { CreateKeyForm } from "../components/create-key-form"; import { KeyRevealCard } from "../components/key-reveal-card"; @@ -166,6 +167,8 @@ export function ApiKeysTab() { )} + + !o && setRevokeTarget(null)}> From e1eaf6e202e5b6b5d7abbb762065a3325b8ee155 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 16 Jun 2026 17:38:42 -0700 Subject: [PATCH 03/26] Studio: HTML canvas cards in chat with auto-render, a Code view, and visible diffusion code (#6374) * Studio: auto-render fenced HTML in chat replies as canvas cards After an assistant reply finishes, append a clickable canvas card for each fenced html block in its text, with no render_html tool call and no extra message. Fragments and non-collapsed documents are covered; full documents already collapsed in place and blocks rendered by the render_html tool are skipped so nothing shows twice. Hoists the fence helpers out of markdown-text.tsx into a shared module (html-fences.ts) and adds a line-based multi-fence scanner so several html blocks in one reply are all found. * Studio: add HTML Code button to canvas cards and keep diffusion code visible When Canvas mode is on or a diffusion model is loaded, the canvas card shows a Preview and an HTML Code button side by side; Code opens the panel source view. The requested view is threaded through openArtifact so the surface opens to preview or code. Diffusion no longer collapses its full HTML answer, so the raw code stays in the message and the trailing canvas card is appended next to it. * Studio: drop the Code button on diffusion cards since their code is already inline * Studio: address review feedback on HTML canvas auto-cards - Build the fence-body indent regex once per fence instead of per line. - Only skip full-doc fences the in-place collapse can render (plain unindented triple-backtick), so 4-backtick or indented docs still get a card. - Scan each text part on its own so a fence cannot stitch across a tool, source, or reasoning part. - Exclude diffusion replies from the collapse/skip gates and the card Code button, since diffusion keeps its HTML inline. --------- Co-authored-by: Daniel Han --- .../components/assistant-ui/markdown-text.tsx | 85 ++-------- .../assistant-ui/message-html-artifacts.tsx | Bin 0 -> 2767 bytes .../src/components/assistant-ui/thread.tsx | 2 + .../features/chat/artifacts/artifact-card.tsx | 106 ++++++++----- .../chat/artifacts/artifact-surface.tsx | 7 + .../features/chat/artifacts/html-fences.ts | 145 ++++++++++++++++++ .../src/features/chat/artifacts/store.ts | 7 +- 7 files changed, 238 insertions(+), 114 deletions(-) create mode 100644 studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx create mode 100644 studio/frontend/src/features/chat/artifacts/html-fences.ts diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index efdef37e72..40fc8b8da6 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -4,6 +4,13 @@ "use client"; import { ArtifactCard, useChatRuntimeStore } from "@/features/chat"; +import { + getCodeFence, + isFullHtmlDocument, + isHtmlFence, + isRenderableRenderHtmlToolPart, + isSvgFence, +} from "@/features/chat/artifacts/html-fences"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { preprocessLaTeX } from "@/lib/latex"; import { openLink } from "@/lib/open-link"; @@ -45,62 +52,16 @@ const STREAMDOWN_COMPONENTS = { }; const COPY_RESET_MS = 2000; const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i; -const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/; const ACTION_PANEL_CLASS = "pointer-events-auto flex shrink-0 items-center gap-1"; const ACTION_BUTTON_CLASS = "flex size-8 cursor-pointer items-center justify-center rounded-[10px] text-chat-icon-fg transition-all hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover disabled:cursor-not-allowed disabled:opacity-50"; -type CodeFence = { - language: string | null; - source: string; -}; - -type ToolCallPartLike = { - type?: string; - toolName?: string; - args?: unknown; - result?: unknown; -}; - -function isRenderableRenderHtmlToolPart(part: unknown): boolean { - const toolPart = part as ToolCallPartLike; - if (toolPart.type !== "tool-call" || toolPart.toolName !== "render_html") { - return false; - } - if ( - typeof toolPart.result === "string" && - toolPart.result.startsWith("Error:") - ) { - return false; - } - if ( - typeof toolPart.result === "string" && - toolPart.result.startsWith("Rendered HTML canvas") - ) { - return true; - } - const args = toolPart.args as { code?: unknown } | undefined; - return typeof args?.code === "string" && args.code.trim().length > 0; -} - function getMermaidSource(blockContent: string): string | null { const source = blockContent.match(MERMAID_SOURCE_RE)?.[1]?.trim(); return source && source.length > 0 ? source : null; } -function getCodeFence(blockContent: string): CodeFence | null { - const match = blockContent.trimEnd().match(CODE_FENCE_RE); - if (!match) { - return null; - } - - return { - language: match[1]?.trim() || null, - source: match[2], - }; -} - function getCodeFilename(language: string | null) { const extByLanguage: Record = { bash: "sh", @@ -131,28 +92,6 @@ function getCodeFilename(language: string | null) { return `snippet.${ext}`; } -function isSvgFence(codeFence: CodeFence): boolean { - const lang = codeFence.language?.toLowerCase() ?? ""; - if (lang === "svg") return true; - if (lang === "xml" || lang === "html") { - const trimmed = codeFence.source.trimStart(); - // Match followed by ]/i.test(trimmed); -} - const UNSAFE_SVG_RE = /]|on\w+\s*=|javascript:|]|]|]|]/i; @@ -285,15 +224,13 @@ function CodeBlockActions({ ); } -// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in -// thread.tsx) and has the HTML canvas feature on by default, so a full-HTML answer -// (e.g. a playable game) renders as an interactive card without the global toggle. +// Collapse a full-HTML answer in place into an artifact card. Diffusion keeps the +// raw code visible instead (the trailing MessageHtmlArtifacts appends its card). function StreamdownBlock(props: BlockProps) { const shouldCollapseHtmlArtifacts = useChatRuntimeStore( (state) => - state.artifactsEnabled || - state.collapseHtmlArtifacts || - state.loadedIsDiffusion, + (state.artifactsEnabled || state.collapseHtmlArtifacts) && + !state.loadedIsDiffusion, ); const messageHasRenderableRenderHtmlTool = useAuiState(({ message }) => message.parts.some(isRenderableRenderHtmlToolPart), diff --git a/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx b/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx new file mode 100644 index 0000000000000000000000000000000000000000..7555287211b21d411d80c88f6b49c092eac99ce0 GIT binary patch literal 2767 zcmb7GT~8x76zy|<#bs2j5;d7@snmyM({0K}7YSHIXjN&oTFi`xadGUC?ExY-|Gnqh zGa-OlsXSx``}%&IbMEy_r?faf_-k^a3TeEY9GB7rRVnG-(fs%GlgUr{+le)LHJDCm z-!?0!Zt6h0Z+Cy5G){V0Q>Z2VVZ64Xrui`)nL*i35{88L$_IgfOXf9hRCMUTa(2k}{-%N1_iQr!$KLO)1@1Hlkpyh6k<3 zWGunLi9Tpf2eOh52Q3%Q{!lmKG$Pb-Roq!sQrR|I6(UIgj{L1^umza<%JvWr0myeB}HgK*tCERSm>P$C}B~Ak3i{&f%Ri7 zDt{5d6W3a%H{QPli08tE@!(%PPq7(u@>52A5}EQ_@i@(Ss}=z;KBc$O!K77FU{!3V zNMY4TI@*sVp30@&oa-$a<2j)+2S!9-y%-zF7Dek=6kKyuMRZXz&$m)GaS0?lfbJ!p zoRhaC*k8;XFp!aVkb(~?6g8+5g|ps+Xk*QUl@=Y)a7$m_wyuz4&@H?*%G`|dL1B## zbUweh{JJ>AV}AMRf@YL`#~mRbBBROXG<5>Q-4hJG9qq$t4Cg6v2q|1MBH1f^88sgaOg; zu|GlUm`JDm93325{~?b&%@9B96JRBc_BQTo(W$D}(u*)JBgC+Ml*m)=?NSc^Tb!d! zRG{>y@(L2(nlu~EJZ6ht(GcTc!_+v=2~3~FAzrK_>6~>wx`oGE*@WkE>9hhER-W!v zs3AF?44iR%0hA8H2`$FY3TL^Ed&oJR-?Y8{KTv?IfKd_QA<8SIgLK0o$BV@9%N83P zBBM8NC`)fhFPiq}m!jD^9Nhpf3>Y@Ys0$8a0_QpZKk85YZIu~57xXjakGiRbUg^f4 zukF2|>s7?$)r8x_AS!gr8m$0&sj8~=${J=9utVuZ9oyE&Az#>1qPMle{h**-4A>`7 z`>E+Ui33~sgzn4v*y)<}D;c=1v|(Jx0b&|uhlZzMN#DK!3z+7;`)55m0>Yhxon)3e z?${q}mx6YHRcg019aC;R*oIFHjtN{53%%<{#D3L1qc{qO2neJ(By>raTOabwuTl^ zVOKvdypqFHzie(k_LTm1Qj}cv1jnnBB_3U3{hR~Dgn6|mtT#{1%)a7uT94`T2I}!^ zpfMFX5XY}weMyWmU8$a2OLK$oCZ2?EdyYCX4D`qxtyu)!T2kTrQr+Pr2{+-HSjt(p z43k}|w0wv^6I4JqqM7Vwy-FXuXt8_a` { /> +
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx index 561cb04041..a1dabe8ecd 100644 --- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx +++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx @@ -3,12 +3,14 @@ "use client"; +import { CodeToggleIcon } from "@/components/assistant-ui/code-toggle-icon"; import { cn } from "@/lib/utils"; import { useAuiState } from "@assistant-ui/react"; import { LayoutTwoColumnIcon as Layout2ColumnIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useLayoutEffect, useMemo } from "react"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import type { ArtifactViewMode } from "./html-frame"; import { hasAutoOpenedArtifact, rememberAutoOpenedArtifact, @@ -20,6 +22,9 @@ import { createChatArtifact, } from "./types"; +const CARD_BASE = + "group/artifact-card relative flex min-h-[52px] cursor-pointer items-center overflow-hidden rounded-lg border border-border/70 bg-muted/15 px-3 py-2 text-left transition-colors hover:bg-muted/25 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 dark:bg-muted/10 dark:hover:bg-muted/20"; + export function ArtifactCard({ code, title, @@ -40,6 +45,11 @@ export function ArtifactCard({ isStreaming?: boolean; }) { const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId); + // Canvas mode collapses the raw code in place, so offer a Code button too. + // Diffusion keeps its code inline, so it needs no Code button. + const showCodeButton = useChatRuntimeStore( + (state) => state.artifactsEnabled && !state.loadedIsDiffusion, + ); const messageIdFromContext = useAuiState(({ message }) => message.id); const threadIdFromContext = useAuiState( ({ threads }) => threads.mainThreadId, @@ -87,7 +97,7 @@ export function ArtifactCard({ } rememberAutoOpenedArtifact(artifact.id); - openArtifact(artifact, { surface }); + openArtifact(artifact, { surface, view: "preview" }); }, [ artifact, autoOpen, @@ -97,45 +107,63 @@ export function ArtifactCard({ updateArtifact, ]); - return ( - +
+ {isCode ? ( + + ) : ( + + )} + + + {isCode ? "HTML Code" : artifact.title} + + + HTML canvas + + + {isStreaming && !isCode ? ( + + Generating + + ) : null} +
+ + ); + }; + + if (!showCodeButton) { + return
{renderButton("preview")}
; + } + + return ( +
+ {renderButton("preview")} + {renderButton("source")} +
); } diff --git a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx index 2d5cbabe8d..b27acdf6ee 100644 --- a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx +++ b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx @@ -31,6 +31,7 @@ import { } from "react"; import { Streamdown } from "streamdown"; import { ArtifactHtmlFrame, type ArtifactViewMode } from "./html-frame"; +import { useChatArtifactsStore } from "./store"; import type { ChatArtifact } from "./types"; import { getArtifactFilename } from "./types"; @@ -119,6 +120,8 @@ export function ArtifactSurface({ onOpenFullscreen?: () => void; }) { const [viewMode, setViewMode] = useState("preview"); + // Follow the view the opener asked for (Preview vs Code button), per artifact. + const requestedView = useChatArtifactsStore((state) => state.requestedView); const [copied, setCopied] = useState(false); const copyResetRef = useRef | null>(null); const surfaceRef = useRef(null); @@ -138,6 +141,10 @@ export function ArtifactSurface({ }; }, []); + useEffect(() => { + setViewMode(requestedView); + }, [artifact.id, requestedView]); + useEffect(() => { if (variant !== "overlay") return; previousFocusRef.current = document.activeElement; diff --git a/studio/frontend/src/features/chat/artifacts/html-fences.ts b/studio/frontend/src/features/chat/artifacts/html-fences.ts new file mode 100644 index 0000000000..5398cf1add --- /dev/null +++ b/studio/frontend/src/features/chat/artifacts/html-fences.ts @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Shared fenced-code helpers for the HTML-artifact render paths, hoisted from +// markdown-text.tsx so the in-place collapse and the post-message auto-render +// agree on what counts as a renderable HTML fence. + +export type CodeFence = { + language: string | null; + source: string; +}; + +// Matches one fenced block spanning the whole string (one pre-split block). +export const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/; + +export type ToolCallPartLike = { + type?: string; + toolName?: string; + args?: unknown; + result?: unknown; +}; + +// True when a part is a render_html tool call with usable code or a non-error result. +export function isRenderableRenderHtmlToolPart(part: unknown): boolean { + const toolPart = part as ToolCallPartLike; + if (toolPart.type !== "tool-call" || toolPart.toolName !== "render_html") { + return false; + } + if ( + typeof toolPart.result === "string" && + toolPart.result.startsWith("Error:") + ) { + return false; + } + if ( + typeof toolPart.result === "string" && + toolPart.result.startsWith("Rendered HTML canvas") + ) { + return true; + } + const args = toolPart.args as { code?: unknown } | undefined; + return typeof args?.code === "string" && args.code.trim().length > 0; +} + +export function getCodeFence(blockContent: string): CodeFence | null { + const match = blockContent.trimEnd().match(CODE_FENCE_RE); + if (!match) { + return null; + } + + return { + language: match[1]?.trim() || null, + source: match[2], + }; +} + +export function isSvgFence(codeFence: CodeFence): boolean { + const lang = codeFence.language?.toLowerCase() ?? ""; + if (lang === "svg") return true; + if (lang === "xml" || lang === "html") { + const trimmed = codeFence.source.trimStart(); + // then ]/i.test(trimmed); +} + +export interface HtmlFence { + source: string; + isFullDocument: boolean; + // Plain 3-backtick unindented fence: the only form the in-place collapser + // (CODE_FENCE_RE) recognizes, so only these may be skipped as already shown. + isPlainFence: boolean; + index: number; +} + +// Opening fence: up to 3 leading spaces, >=3 backticks, then a backtick-free info string. +const FENCE_OPEN_RE = /^( {0,3})(`{3,})([^`\r\n]*)$/; + +// Scan a full message for every closed ```html fence. Line-based so multiple +// fences are found and backticks inside a {label} @@ -293,6 +293,7 @@ export function AppSidebar() { }; const isRecipesRoute = pathname.startsWith("/data-recipes"); + const isExportRoute = pathname === "/export" || pathname.startsWith("/export/"); const { displayTitle, avatarDataUrl } = useEffectiveProfile(); const { projects } = useChatProjects(); @@ -334,10 +335,14 @@ export function AppSidebar() { undefined : undefined; - // Training runs + // Training runs: surfaced as sidebar "Recents" on Train, Recipes, and Export, + // falling back to chat recents when there are no runs yet. + const trainingRecentsRoute = isStudioRoute || isRecipesRoute || isExportRoute; const { items: runItems } = useTrainingHistorySidebarItems( - !chatOnly && isStudioRoute, + !chatOnly && trainingRecentsRoute, ); + const showTrainingRecents = + !chatOnly && trainingRecentsRoute && runItems.length > 0; const activeJobId = useTrainingRuntimeStore((s) => s.jobId); const currentRunViewActive = useTrainingRuntimeStore((s) => s.currentRunViewActive); const selectedHistoryRunId = useTrainingRuntimeStore((s) => s.selectedHistoryRunId); @@ -667,7 +672,7 @@ export function AppSidebar() { ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; const buttonClass = cn( - "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium", + "sidebar-nav-btn h-[33px] cursor-pointer rounded-[14px] pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium", // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the // title with the nav items above. variant === "project" ? "pl-[39px]" : "pl-3", @@ -921,7 +926,7 @@ export function AppSidebar() { @@ -1278,14 +1286,14 @@ export function AppSidebar() {
diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 3eaa21a19e..a6c326c64e 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -20,7 +20,9 @@ import { } from "@assistant-ui/react"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { type VariantProps, cva } from "class-variance-authority"; -import { CheckIcon, ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react"; +import { ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { HugeiconsIcon } from "@hugeicons/react"; import { type CSSProperties, type ComponentProps, @@ -293,7 +295,7 @@ function ReasoningCopyButton({ startIndex, endIndex }: { startIndex: number; end aria-label="Copy reasoning" > {copied ? ( - + ) : ( )} diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index 965ad56bfa..20a22c3a4a 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -17,11 +17,12 @@ import { } from "@assistant-ui/react"; import { AlertCircleIcon, - CheckIcon, ChevronDownIcon, LoaderIcon, XCircleIcon, } from "lucide-react"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { HugeiconsIcon } from "@hugeicons/react"; import { type CSSProperties, type ComponentProps, @@ -95,9 +96,15 @@ function ToolFallbackRoot({ type ToolStatus = ToolCallMessagePartStatus["type"]; +// The shared app tick is icon data, not a component; wrap it to slot into the +// status map alongside the lucide icons. +function CompleteTickIcon(props: Omit, "icon">) { + return ; +} + const statusIconMap: Record = { running: LoaderIcon, - complete: CheckIcon, + complete: CompleteTickIcon, incomplete: XCircleIcon, "requires-action": AlertCircleIcon, }; diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx index 23a83bc5ba..84309815e5 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx @@ -8,12 +8,9 @@ import { type ToolCallMessagePartComponent, useAuiState, } from "@assistant-ui/react"; -import { - CheckIcon, - CopyIcon, - FileTextIcon, - TerminalIcon, -} from "lucide-react"; +import { CopyIcon, FileTextIcon, TerminalIcon } from "lucide-react"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { HugeiconsIcon } from "@hugeicons/react"; import { Spinner } from "@/components/ui/spinner"; import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react"; import { @@ -93,7 +90,7 @@ function CopyBtn({ text }: { text: string }) { aria-label="Copy to clipboard" > {copied ? ( - + ) : ( )} diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx index d7b0e58f18..1afffaed7d 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx @@ -7,7 +7,9 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { getAuthToken } from "@/features/auth/session"; import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; import { code as codePlugin } from "@streamdown/code"; -import { CheckIcon, CodeIcon, CopyIcon } from "lucide-react"; +import { CodeIcon, CopyIcon } from "lucide-react"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { HugeiconsIcon } from "@hugeicons/react"; import { Spinner } from "@/components/ui/spinner"; import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react"; import { Streamdown } from "streamdown"; @@ -63,7 +65,7 @@ function CopyBtn({ text }: { text: string }) { aria-label="Copy to clipboard" > {copied ? ( - + ) : ( )} diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx index 3c12b16e6b..c088d5cf98 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx @@ -5,7 +5,9 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard"; import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; -import { CheckIcon, CopyIcon, TerminalIcon } from "lucide-react"; +import { CopyIcon, TerminalIcon } from "lucide-react"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { HugeiconsIcon } from "@hugeicons/react"; import { Spinner } from "@/components/ui/spinner"; import { memo, useCallback, useEffect, useRef, useState } from "react"; import { @@ -53,7 +55,7 @@ function CopyBtn({ text }: { text: string }) { aria-label="Copy to clipboard" > {copied ? ( - + ) : ( )} diff --git a/studio/frontend/src/components/ui/combobox.tsx b/studio/frontend/src/components/ui/combobox.tsx index 7bf0d7bf75..db2d494b6d 100644 --- a/studio/frontend/src/components/ui/combobox.tsx +++ b/studio/frontend/src/components/ui/combobox.tsx @@ -7,7 +7,7 @@ import { Combobox as ComboboxPrimitive } from "@base-ui/react"; import * as React from "react"; -import { createContext, useContext, useState } from "react"; +import { createContext, useState } from "react"; import { Button } from "@/components/ui/button"; import { useDialogPortalContainer } from "@/components/ui/dialog"; @@ -18,11 +18,9 @@ import { InputGroupInput, } from "@/components/ui/input-group"; import { Tick02Icon } from "@/lib/tick-icon"; +import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { cn } from "@/lib/utils"; -import { - ArrowDown01Icon, - Cancel01Icon, -} from "@hugeicons/core-free-icons"; +import { Cancel01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; const ComboboxOpenContext = createContext(false); @@ -68,7 +66,7 @@ function ComboboxTrigger({ > {children} @@ -107,18 +105,8 @@ function ComboboxInput({ showTrigger?: boolean; showClear?: boolean; }): React.ReactElement { - const isOpen = useContext(ComboboxOpenContext); - return ( - + } {...props} @@ -130,7 +118,7 @@ function ComboboxInput({ variant="ghost" asChild data-slot="input-group-button" - className="group-has-data-[slot=combobox-clear]/input-group:hidden data-pressed:bg-transparent" + className="group-has-data-[slot=combobox-clear]/input-group:hidden bg-transparent hover:bg-transparent data-pressed:bg-transparent aria-expanded:bg-transparent dark:hover:bg-transparent" disabled={disabled} > @@ -209,7 +197,7 @@ function ComboboxItem({ {copied ? ( - + ) : ( )} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index f3cbd2e966..3b518f83bf 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -2152,16 +2152,6 @@ export function ChatPage(): ReactElement { className="max-w-[62vw] !pr-3 sm:max-w-none !h-[34px]" /> )} - {incognito && view.mode === "single" && ( -
- - Temporary -
- )} {view.mode !== "compare" && currentProjectId && (
@@ -1387,7 +1387,7 @@ function NewPromptForm({ onClose, onRefresh }: { onClose: () => void; onRefresh: Cancel @@ -1486,7 +1486,7 @@ function PromptListCard({ onClick={handleSave} disabled={items.filter((t) => t.trim()).length === 0} > - Save List + Save List @@ -1637,7 +1637,7 @@ function NewPromptListForm({ onClose, onRefresh }: { onClose: () => void; onRefr onClick={handleSave} disabled={items.filter((t) => t.trim()).length === 0} > - Save Prompt List + Save Prompt List diff --git a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx index 76f443afe4..005455cd94 100644 --- a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx +++ b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx @@ -41,6 +41,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useNavigate } from "@tanstack/react-router"; +import { RecentTrainingsSection } from "@/features/studio/recent-trainings-section"; import type { ReactElement } from "react"; import { useEffect, useState } from "react"; import { @@ -528,6 +529,8 @@ export function DataRecipesPage(): ReactElement { ))} )} + + diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 8180849733..e467b10d24 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { SectionCard } from "@/components/section-card"; +import { RecentTrainingsSection } from "@/features/studio/recent-trainings-section"; import { Button } from "@/components/ui/button"; import { Combobox, @@ -1129,6 +1130,8 @@ export function ExportPage() { )} + + void; +}) { + return ( + + ); +} + +export function CardCarousel({ + items, + getKey, + renderItem, + itemWidth, + itemHeight, + ariaLabel, +}: { + items: T[]; + getKey: (item: T) => string; + renderItem: (item: T) => ReactNode; + itemWidth: number; + itemHeight: number; + ariaLabel: string; +}) { + const scrollerRef = useRef(null); + const [canLeft, setCanLeft] = useState(false); + const [canRight, setCanRight] = useState(false); + const stepPx = itemWidth + CARD_GAP_PX; + const arrowCenterPx = CAROUSEL_TOP_PADDING_PX + itemHeight / 2; + + const updateArrows = useCallback(() => { + const el = scrollerRef.current; + if (!el) return; + setCanLeft(el.scrollLeft > 1); + setCanRight(el.scrollLeft + el.clientWidth < el.scrollWidth - 1); + }, []); + + useEffect(() => { + const el = scrollerRef.current; + if (!el) return; + updateArrows(); + const observer = new ResizeObserver(updateArrows); + observer.observe(el); + return () => observer.disconnect(); + }, [updateArrows]); + + useEffect(() => { + updateArrows(); + }, [updateArrows, items]); + + const scrollByCards = useCallback( + (direction: 1 | -1) => { + scrollerRef.current?.scrollBy({ + left: direction * stepPx, + behavior: "smooth", + }); + }, + [stepPx], + ); + + return ( +
+
+ {items.map((item) => ( +
+ {renderItem(item)} +
+ ))} +
+ - {/* Train CTA hidden until Hub->train picker ships; divider pairs with it. */} + {/* Train CTA hidden until Hub→train picker ships; divider pairs with it. */} {(!isDownloaded || downloading || HUB_POST_DOWNLOAD_ACTIONS_VISIBLE) && ( )} diff --git a/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx b/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx index d1b8671bfe..48bb9c1bce 100644 --- a/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx +++ b/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx @@ -8,7 +8,8 @@ import { HugeiconsIcon } from "@hugeicons/react"; /** * Inspector action-button affordance during a download: spinner that cross-fades * to a cancel glyph on `.hub-action-btn` hover, in the same 16x16 slot so the - * percentage label never shifts. The swap is pure CSS; the component only carries + * percentage label never shifts. The swap is pure CSS + * (`.hub-action-btn:hover .hub-cta-indicator-*`); the component only carries * the marker classes. */ export function DownloadCancelIndicator() { diff --git a/studio/frontend/src/features/hub/catalog/download-card.tsx b/studio/frontend/src/features/hub/catalog/download-card.tsx index e5253526ab..5083fb2f84 100644 --- a/studio/frontend/src/features/hub/catalog/download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/download-card.tsx @@ -34,8 +34,9 @@ import { } from "./use-download-card-state"; /** - * Shared shell for every download surface (safetensors, GGUF, dataset): card frame, - * progress bar, transport-conflict dialog, plus card-specific `dialogs` and children. + * Shared shell for every download surface (safetensors, GGUF, dataset): card + * frame, progress bar, transport-conflict dialog, plus card-specific `dialogs` + * and children. */ export function DownloadCard({ job, @@ -94,7 +95,7 @@ export function CardDeleteButton({ e.stopPropagation(); onClick(); }} - className="inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-[8px] text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-rose-500/10 hover:text-rose-600 focus-visible:opacity-100 group-hover/dl:opacity-100 dark:hover:bg-rose-500/15 dark:hover:text-rose-400" + className="inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-rose-500/10 hover:text-rose-600 focus-visible:opacity-100 group-hover/dl:opacity-100 dark:hover:bg-rose-500/15 dark:hover:text-rose-400" > - + Open external link diff --git a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx index 8633bda6fa..ea3f158aa6 100644 --- a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx @@ -26,7 +26,6 @@ import { normalizeGgufVariantIdentity, } from "../lib/model-identity"; import { cn } from "@/lib/utils"; -import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { useHfTokenStore } from "../stores/hf-token-store"; import { Delete02Icon, @@ -35,6 +34,7 @@ import { PencilEdit02Icon, PlayIcon, } from "@hugeicons/core-free-icons"; +import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { memo, @@ -588,7 +588,7 @@ export function GgufDownloadCard({ const variantListUnavailable = !sortedVariants || sortedVariants.length === 0; const showVariantLoadingState = loading && variantListUnavailable; - // Keep showing download progress even when the variant list is unavailable, so a + // Keep showing download progress while the variant list is unavailable, so a // remount never hides an in-flight download behind the variant status card. if (progress && variantListUnavailable) { return ( @@ -666,7 +666,7 @@ export function GgufDownloadCard({ e.preventDefault(); setOpen((o) => !o); }} - className="hub-menu-trigger flex h-9 min-w-0 flex-1 cursor-pointer items-center gap-2.5 rounded-full px-3 text-left transition-colors hover:bg-foreground/[0.04] data-[state=open]:bg-foreground/[0.06] dark:hover:bg-white/[0.1] dark:data-[state=open]:bg-white/[0.06]" + className="hub-menu-trigger flex h-9 min-w-0 flex-1 cursor-pointer items-center gap-2.5 rounded-full px-3 text-left transition-colors hover:bg-foreground/[0.04] data-[state=open]:bg-foreground/[0.06] dark:hover:bg-white/[0.04] dark:data-[state=open]:bg-white/[0.06]" > {selected ? ( @@ -723,7 +722,7 @@ export function GgufDownloadCard({ @@ -831,11 +830,7 @@ export function GgufDownloadCard({ ) : selected?.downloaded ? ( <> - + Run ) : ( diff --git a/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx new file mode 100644 index 0000000000..460dba9613 --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx @@ -0,0 +1,81 @@ +// 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 { ArrowLeft01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { cn } from "@/lib/utils"; +import { type ComponentProps, useEffect, useRef, useState } from "react"; +import { ModelInspector } from "./model-inspector"; + +type InspectorProps = ComponentProps; + +export function HubDetailView({ + onBack, + compact = false, + ...inspectorProps +}: InspectorProps & { onBack: () => void; compact?: boolean }) { + const scrollRef = useRef(null); + const [scrolled, setScrolled] = useState(false); + // Split pane is narrower than the full-page overlay; tighter measure reads better. + const measure = compact + ? "mx-auto w-full max-w-[860px] px-5 sm:px-5" + : "mx-auto w-full max-w-[1100px] px-5 sm:px-8"; + + useEffect(() => { + const el = scrollRef.current; + if (!el) return; + const onScroll = () => { + const next = el.scrollTop > 0; + setScrolled((current) => (current === next ? current : next)); + }; + onScroll(); + el.addEventListener("scroll", onScroll, { passive: true }); + return () => el.removeEventListener("scroll", onScroll); + }, []); + + return ( +
+ {/* Same top scroll fade as the left column. The sticky back-bar, when + shown, sits above and hides it. */} + + ); +} diff --git a/studio/frontend/src/features/hub/catalog/hub-feed.tsx b/studio/frontend/src/features/hub/catalog/hub-feed.tsx new file mode 100644 index 0000000000..db39e31b7b --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/hub-feed.tsx @@ -0,0 +1,40 @@ +// 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 { memo } from "react"; +import { HUB_SECTION_TITLE, type HubSection } from "../lib/channels"; +import type { DiscoverRow } from "../types"; +import { HubSectionRow } from "./hub-section-row"; + +export interface HubFeedSectionData { + rows: DiscoverRow[]; + isLoading: boolean; +} + +export const HubFeed = memo(function HubFeed({ + trending, + deviceType, + isDataset, + onSelect, + onOpenChannel, +}: { + trending: HubFeedSectionData; + deviceType: string | null; + isDataset: boolean; + onSelect: (id: string) => void; + onOpenChannel: (section: HubSection) => void; +}) { + return ( +
+ onOpenChannel("trending")} + deviceType={deviceType} + isDataset={isDataset} + /> +
+ ); +}); diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx index db1332f03a..ed83e46b36 100644 --- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx @@ -6,9 +6,9 @@ import { PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; +import { cn } from "@/lib/utils"; import { Tick02Icon } from "@/lib/tick-icon"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; -import { cn } from "@/lib/utils"; import { HugeiconsIcon } from "@hugeicons/react"; import { type KeyboardEvent, @@ -167,18 +167,13 @@ export function HubOptionMenu({ } }} > - - {triggerContent ?? ( - - {selected?.triggerLabel ?? selected?.label ?? value} - - )} + + {triggerContent ?? selected?.triggerLabel ?? selected?.label ?? value} {showChevron && ( )} @@ -186,11 +181,11 @@ export function HubOptionMenu({ event.preventDefault()} className={cn( - "hub-menu-instant menu-soft-surface w-max min-w-[var(--radix-popover-trigger-width)] max-w-[min(var(--radix-popover-content-available-width),calc(100vw-1rem))] rounded-[21px] px-[9px] py-2 ring-0", + "hub-menu-instant menu-soft-surface w-max min-w-[var(--radix-popover-trigger-width)] max-w-[min(var(--radix-popover-content-available-width),calc(100vw-1rem))] rounded-[14px] p-1 ring-0", contentClassName, )} > @@ -219,14 +214,14 @@ export function HubOptionMenu({ }} onPointerEnter={() => activateIndex(index)} className={cn( - "relative flex w-full min-w-0 cursor-pointer select-none items-center gap-2.5 rounded-[12px] py-2 px-3 text-left text-sm leading-snug outline-none transition-colors", + "relative flex w-full min-w-0 cursor-pointer select-none items-center rounded-[12px] py-2 pr-8 pl-3 text-left text-sm leading-snug outline-none transition-colors", )} > {option.label} {selectedOption && ( - + + {SKELETON_KEYS.map((key) => ( + + ))} +
+ ); +} + +export const HubSectionRow = memo(function HubSectionRow({ + title, + rows, + onSelect, + onOpenList, + deviceType, + isDataset, + isLoading, +}: { + title: string; + rows: DiscoverRow[]; + onSelect: (id: string) => void; + onOpenList: () => void; + deviceType: string | null; + isDataset: boolean; + isLoading: boolean; +}) { + const showSkeleton = isLoading && rows.length === 0; + if (!showSkeleton && rows.length === 0) { + return null; + } + + return ( +
+

+ +

+ {showSkeleton ? ( + + ) : ( + row.id} + itemWidth={MODEL_CARD_WIDTH_PX} + itemHeight={MODEL_CARD_HEIGHT_PX} + ariaLabel={title} + renderItem={(row) => ( + + )} + /> + )} +
+ ); +}); diff --git a/studio/frontend/src/features/hub/catalog/hub-top-bar.tsx b/studio/frontend/src/features/hub/catalog/hub-top-bar.tsx new file mode 100644 index 0000000000..f706dd1ff2 --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/hub-top-bar.tsx @@ -0,0 +1,14 @@ +// 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 type { ReactNode } from "react"; + +export function HubTopBar({ children }: { children: ReactNode }) { + return ( +
+
+ {children} +
+
+ ); +} diff --git a/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx b/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx index 5b9ca6b1de..ddda124745 100644 --- a/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx +++ b/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx @@ -47,7 +47,7 @@ export function LocalDatasetCard({ className="ml-1 mr-0 h-5 w-px shrink-0 bg-foreground/[0.06] opacity-100 transition-opacity duration-150 group-hover/dl:opacity-0 dark:bg-white/[0.04]" /> )} - {/* Train CTA hidden until Hub->train picker ships. */} + {/* Train CTA hidden until Hub→train picker ships. */} {onTrain && HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && ( @@ -358,7 +357,7 @@ export function LocalOnDeviceCard({ @@ -452,7 +451,7 @@ export function LocalOnDeviceCard({ !runActionsVisible && "hidden", )} > - {onTrain && ( + {onTrain && HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && ( + ); + + if (!tip) { + return card; + } + return ( + + {card} + + {tip} + + + ); +}); diff --git a/studio/frontend/src/features/hub/catalog/model-inspector.tsx b/studio/frontend/src/features/hub/catalog/model-inspector.tsx index 24909b38aa..371d595ddc 100644 --- a/studio/frontend/src/features/hub/catalog/model-inspector.tsx +++ b/studio/frontend/src/features/hub/catalog/model-inspector.tsx @@ -12,13 +12,17 @@ import { classifyUnslothSupport, } from "@/features/hub/hooks/use-hub-model-search"; import { useOnlineStatus } from "@/features/hub/hooks/use-online-status"; -import { formatBytes, formatRelativeShort } from "@/features/hub/lib/format"; -import { Tick02Icon } from "@/lib/tick-icon"; +import { + formatBytes, + formatRelativeShort, + formatShortDate, +} from "@/features/hub/lib/format"; import { cn, formatCompact } from "@/lib/utils"; import { confirmExternalLink } from "../stores/external-link-confirm"; import { useHfTokenStore } from "@/features/hub/stores/hf-token-store"; import { Calendar03Icon, + CalendarAdd01Icon, Copy01Icon, CpuIcon, CubeIcon, @@ -27,40 +31,33 @@ import { FavouriteIcon, Globe02Icon, LayersLogoIcon, + LibraryIcon, LicenseIcon, PackageIcon, RamMemoryIcon, Share05Icon, } from "@hugeicons/core-free-icons"; +import { Tick02Icon } from "@/lib/tick-icon"; import type { IconSvgElement } from "@hugeicons/react"; import { HugeiconsIcon } from "@hugeicons/react"; -import { - memo, - useDeferredValue, - useEffect, - useMemo, - useRef, - useState, -} from "react"; +import { memo, useDeferredValue, useMemo } from "react"; import { useCopyFeedback } from "../hooks/use-copy-feedback"; import { useDatasetSize } from "../hooks/use-dataset-size"; import { + formatLibrary, formatLocalUpdated, formatPipelineTag, parseLanguageTags, } from "../lib/view-models"; import type { SelectedModelView } from "../types"; -import { - selectActiveJob, - useDownloadManagerStore, -} from "../download-manager"; +import { selectActiveJob, useDownloadManagerStore } from "../download-manager"; import { DatasetDownloadSection } from "./dataset-download-section"; import { DownloadSection } from "./download-section"; import { LocalDatasetCard } from "./local-dataset-card"; import { LocalOnDeviceCard } from "./local-on-device-card"; import { ModelReadme } from "./model-readme"; import { OwnerAvatar } from "./owner-avatar"; -import { CapabilityPill } from "./shared"; +import { AccessChip, CapabilityPill } from "./shared"; function ViewRepositoryButton({ repoId, @@ -72,9 +69,13 @@ function ViewRepositoryButton({ const online = useOnlineStatus(); const url = `https://huggingface.co/${isDataset ? "datasets/" : ""}${repoId}`; const baseClass = - "inline-flex size-7 shrink-0 items-center justify-center rounded-[8px] text-muted-foreground transition-colors"; + "inline-flex size-6 shrink-0 items-center justify-center rounded-full text-muted-foreground transition-colors"; const icon = ( - + ); return ( @@ -128,12 +129,12 @@ function CopyRepoButton({ repoId }: { repoId: string }) { type="button" aria-label="Copy repository ID" onClick={handleCopy} - className="inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-[8px] text-muted-foreground transition-colors hover:bg-muted hover:text-foreground" + className="inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground" > @@ -148,10 +149,12 @@ function StatRow({ label, value, icon, + tooltip, }: { label: string; value: string; icon: IconSvgElement; + tooltip?: React.ReactNode; }) { return ( @@ -165,7 +168,9 @@ function StatRow({ {value} - {label} + + {tooltip ?? label} + ); } @@ -175,11 +180,7 @@ function StatGrid({ children }: { children: React.ReactNode }) { } function InspectorDownloadSlot({ children }: { children: React.ReactNode }) { - return ( -
- {children} -
- ); + return
{children}
; } function StatusChip({ @@ -200,7 +201,7 @@ function StatusChip({ return ( void; +}) { + const content = ( + <> + + Base + + {baseModel} + + + ); + + return ( + + + {onSearchHub ? ( + + ) : ( + + {content} + + )} + + + Search this base model in Hub + + + ); +} + type VramInfo = { est: number; status: "fits" | "tight" | "exceeds" } | null; function ModelStatusChips({ @@ -270,8 +318,7 @@ function ModelStatusChips({ )} - Still downloadable to your Hugging Face cache, shared with every - framework that reads it. + Still downloadable to your Hugging Face cache. @@ -322,6 +369,7 @@ export type ModelInspectorActions = { onUseInChat: () => void; onTrain?: () => void; onInventoryChange?: () => void; + onSearchHub?: (query: string) => void; }; export const ModelInspector = memo(function ModelInspector({ @@ -349,17 +397,23 @@ export const ModelInspector = memo(function ModelInspector({ gpuGb, systemRamGb, } = runtime; - const { onLoad, onLoadLocal, onUseInChat, onTrain, onInventoryChange } = - actions; + const { + onLoad, + onLoadLocal, + onUseInChat, + onTrain, + onInventoryChange, + onSearchHub, + } = actions; const deviceType = usePlatformStore((s) => s.deviceType); const hfToken = useHfTokenStore((s) => s.token); const datasetRepoId = isDataset && model?.hubRepoId ? model.hubRepoId : null; const datasetSize = useDatasetSize(datasetRepoId, { token: hfToken || undefined, }); - // Inventory rows are snapshots; the download manager is the live source of truth. - // When a download is in flight, route through the download-aware section so - // progress/cancel stays visible across refreshes. + // Inventory rows are snapshots; the download manager is the live source of + // truth. When a download is in flight, route through the download-aware section + // so progress/cancel stays visible across refreshes. const activeDownloadRepoId = model?.hubRepoId ?? null; const hasActiveHubDownload = useDownloadManagerStore((state) => activeDownloadRepoId @@ -403,35 +457,12 @@ export const ModelInspector = memo(function ModelInspector({ supportTagsKey, ]); - const descScrollRef = useRef(null); - const descScrollKey = `${model?.id ?? ""}\0${readmeRepoId ?? ""}`; const deferredReadmeRepoId = useDeferredValue(readmeRepoId); const readmeReady = deferredReadmeRepoId === readmeRepoId; - const [descScrollState, setDescScrollState] = useState({ - key: descScrollKey, - scrolled: false, - }); - const descScrolled = - descScrollState.key === descScrollKey && descScrollState.scrolled; - useEffect(() => { - const el = descScrollRef.current; - if (!el) return; - el.scrollTop = 0; - const onScroll = () => { - const scrolled = el.scrollTop > 0; - setDescScrollState((current) => - current.key === descScrollKey && current.scrolled === scrolled - ? current - : { key: descScrollKey, scrolled }, - ); - }; - el.addEventListener("scroll", onScroll, { passive: true }); - return () => el.removeEventListener("scroll", onScroll); - }, [descScrollKey]); if (!model) { return ( -
+
@@ -453,8 +484,23 @@ export const ModelInspector = memo(function ModelInspector({ ? formatRelativeShort(model.updatedAt) : formatLocalUpdated(model.localUpdatedAt); const updatedLabel = updatedRaw === "Unknown update" ? "N/A" : updatedRaw; + const createdLabel = model.createdAt ? formatShortDate(model.createdAt) : null; + const libraryLabel = isDataset ? null : formatLibrary(model.libraryName); + const gatedAccess = model.gated !== false && model.gated !== undefined; + const downloadsTooltip = + model.downloadsAllTime != null ? ( + <> + Downloads (30 days) + + {formatCompact(model.downloadsAllTime)} all time + + + ) : ( + "Downloads" + ); const taskLabel = formatPipelineTag(model.pipelineTag) ?? "General"; const licenseLabel = model.license ?? "N/A"; + const baseModelSearchTerm = model.baseModelHubId ?? model.baseModel ?? null; const paramsLabel = model.totalParams ? formatCompact(model.totalParams) : "N/A"; @@ -473,17 +519,17 @@ export const ModelInspector = memo(function ModelInspector({ datasetSize?.numBytesParquet ?? datasetSize?.numBytesOriginal ?? null; return ( -
-
-
+
+
+
-

+

{model.title}

{model.hubRepoId && ( @@ -496,12 +542,12 @@ export const ModelInspector = memo(function ModelInspector({
)}
-
+
{model.owner} {model.owner.toLowerCase() === "unsloth" && ( )}
@@ -524,9 +570,18 @@ export const ModelInspector = memo(function ModelInspector({ {taskLabel} )} + {model.private && } + {gatedAccess && } {model.capabilities.map((capability) => ( ))} + {!isDataset && model.baseModel && baseModelSearchTerm && ( + + )}
@@ -620,7 +675,7 @@ export const ModelInspector = memo(function ModelInspector({ )} -
+
{selectionHiddenByFilters && (

Current selection is hidden by the active filters or search. @@ -634,8 +689,16 @@ export const ModelInspector = memo(function ModelInspector({ )} + {createdLabel && ( + + )} )} + {libraryLabel && ( + + )} {!isDataset && ( )} - {isDataset && languages.length > 0 && ( + {languages.length > 0 && (

-