onSelectRun(run.id)}
onKeyDown={(e) => {
@@ -411,6 +438,38 @@ export function HistoryCardGrid({
{isResuming ? t("studio.history.resuming") : t("studio.history.resumeTraining")}
)}
+ {canCopyPreview && (
+
{run.loss_sparkline && run.loss_sparkline.length >= 2 && (
-
+
Date: Wed, 24 Jun 2026 06:37:41 -0700
Subject: [PATCH 54/56] Verify DiffusionGemma visual-server binary against
approved checksums (#6635)
ensure_diffusion_visual_server() downloaded the visual-server release
asset with the unverified download_file() and marked it executable,
bypassing the approved-checksum manifest that gates every other prebuilt
llama.cpp artifact. The backend later auto-discovers that binary and
launches it through DG_VISUAL_BIN, so a compromised or substituted
release asset could place attacker-controlled native code in the install
tree and have it executed under the Studio user.
Require the matched asset to be present in the approved checksum manifest
and download it through download_file_verified() with the published
sha256. A name-matching asset that is absent from the manifest is refused
rather than executed.
Add regression tests covering the verified-download path and the refusal
of an unapproved asset.
---
studio/install_llama_prebuilt.py | 44 +++++--
.../test_install_llama_prebuilt_logic.py | 119 ++++++++++++++++++
2 files changed, 154 insertions(+), 9 deletions(-)
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index 33cb709ba7..c7f34e39f2 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -4261,7 +4261,10 @@ def ensure_converter_scripts(install_dir: Path, llama_tag: str) -> None:
def ensure_diffusion_visual_server(
- install_dir: Path, host: HostInfo, release_tag: str | None
+ install_dir: Path,
+ host: HostInfo,
+ release_tag: str | None,
+ approved_checksums: ApprovedReleaseChecksums,
) -> None:
"""Best-effort placement of the DiffusionGemma visual-server binary next to
llama-server in the install tree, so Studio can serve DiffusionGemma GGUFs
@@ -4293,6 +4296,7 @@ def ensure_diffusion_visual_server(
try:
assets = github_release_assets(DEFAULT_PUBLISHED_REPO, release_tag)
match = None
+ unapproved_matches: list[str] = []
for asset_name, url in assets.items():
low = asset_name.lower()
if "llama-diffusion-gemma-visual-server" not in low:
@@ -4301,19 +4305,39 @@ def ensure_diffusion_visual_server(
continue
if (not host.is_windows) and low.endswith(".exe"):
continue
- match = (asset_name, url)
+ # This binary is chmod'd executable and later launched by the
+ # backend, so it must be covered by the approved checksum manifest
+ # just like every other prebuilt artifact. An asset that matches the
+ # name but is missing from the manifest is refused rather than run.
+ approved = approved_checksums.artifacts.get(asset_name)
+ if approved is None:
+ unapproved_matches.append(asset_name)
+ continue
+ match = (asset_name, url, approved.sha256)
break
if match is None:
- log(
- "diffusion visual server not found in the published release; native "
- "DiffusionGemma serving needs DG_VISUAL_BIN or a source build"
- )
+ if unapproved_matches:
+ log(
+ "diffusion visual server asset(s) were present but omitted from the "
+ "approved checksum manifest; refusing unverified native executable: "
+ + ", ".join(unapproved_matches)
+ )
+ else:
+ log(
+ "diffusion visual server not found in the published release; native "
+ "DiffusionGemma serving needs DG_VISUAL_BIN or a source build"
+ )
return
bin_dir.mkdir(parents = True, exist_ok = True)
- download_file(match[1], target)
+ download_file_verified(
+ match[1],
+ target,
+ expected_sha256 = match[2],
+ label = f"diffusion visual server {match[0]}",
+ )
if not host.is_windows:
target.chmod(0o755)
- log(f"installed diffusion visual server: {match[0]}")
+ log(f"installed verified diffusion visual server: {match[0]}")
except Exception as exc:
log(
"diffusion visual server fetch skipped "
@@ -6637,7 +6661,9 @@ def install_prebuilt(
f"({textwrap.shorten(str(exc), width = 200, placeholder = '...')})"
)
try:
- ensure_diffusion_visual_server(install_dir, host, plan.release_tag)
+ ensure_diffusion_visual_server(
+ install_dir, host, plan.release_tag, plan.approved_checksums
+ )
except Exception as exc:
log(
"diffusion visual server step skipped; install remains valid "
diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py
index a5a1131ecc..9ee8759bb4 100644
--- a/tests/studio/install/test_install_llama_prebuilt_logic.py
+++ b/tests/studio/install/test_install_llama_prebuilt_logic.py
@@ -37,6 +37,41 @@ install_prebuilt = INSTALL_LLAMA_PREBUILT.install_prebuilt
write_prebuilt_metadata = INSTALL_LLAMA_PREBUILT.write_prebuilt_metadata
existing_install_matches_plan = INSTALL_LLAMA_PREBUILT.existing_install_matches_plan
existing_install_matches_choice = INSTALL_LLAMA_PREBUILT.existing_install_matches_choice
+ensure_diffusion_visual_server = INSTALL_LLAMA_PREBUILT.ensure_diffusion_visual_server
+
+
+def linux_host() -> HostInfo:
+ return HostInfo(
+ system = "Linux",
+ machine = "x86_64",
+ is_windows = False,
+ is_linux = True,
+ is_macos = False,
+ is_x86_64 = True,
+ is_arm64 = False,
+ nvidia_smi = None,
+ driver_cuda_version = None,
+ compute_caps = [],
+ visible_cuda_devices = None,
+ has_physical_nvidia = False,
+ has_usable_nvidia = False,
+ )
+
+
+def approved_release_checksums_for_asset(asset_name: str, sha256: str) -> ApprovedReleaseChecksums:
+ return ApprovedReleaseChecksums(
+ repo = "unslothai/llama.cpp",
+ release_tag = "b9334",
+ upstream_tag = "b9334",
+ artifacts = {
+ asset_name: ApprovedArtifactHash(
+ asset_name = asset_name,
+ sha256 = sha256,
+ repo = "unslothai/llama.cpp",
+ kind = "diffusion-visual-server",
+ )
+ },
+ )
def approved_checksums_for(
@@ -2828,3 +2863,87 @@ def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", True)
calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32)
assert calls == {"quantize": 1, "server": 1}
+
+
+def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tmp_path: Path):
+ asset_name = "llama-diffusion-gemma-visual-server-linux-x64"
+ expected_sha = "a" * 64
+ asset_url = "https://github.com/unslothai/llama.cpp/releases/download/b9334/" + asset_name
+ calls: list[tuple[str, Path, str | None, str | None]] = []
+
+ monkeypatch.setattr(
+ INSTALL_LLAMA_PREBUILT,
+ "github_release_assets",
+ lambda repo, tag: {asset_name: asset_url},
+ )
+
+ def fake_download_file(url, destination):
+ raise AssertionError("diffusion visual server must not use unverified download_file")
+
+ def fake_download_file_verified(url, destination, *, expected_sha256, label):
+ calls.append((url, Path(destination), expected_sha256, label))
+ Path(destination).write_bytes(b"verified visual server")
+
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+ monkeypatch.setattr(
+ INSTALL_LLAMA_PREBUILT, "download_file_verified", fake_download_file_verified
+ )
+
+ ensure_diffusion_visual_server(
+ tmp_path / "install",
+ linux_host(),
+ "b9334",
+ approved_release_checksums_for_asset(asset_name, expected_sha),
+ )
+
+ target = tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server"
+ assert calls == [
+ (
+ asset_url,
+ target,
+ expected_sha,
+ f"diffusion visual server {asset_name}",
+ )
+ ]
+ assert target.read_bytes() == b"verified visual server"
+ assert target.stat().st_mode & 0o777 == 0o755
+
+
+def test_diffusion_visual_server_refuses_unapproved_release_asset(monkeypatch, tmp_path: Path):
+ asset_name = "llama-diffusion-gemma-visual-server-attacker-linux"
+ verified_calls: list[str] = []
+ raw_calls: list[str] = []
+
+ monkeypatch.setattr(
+ INSTALL_LLAMA_PREBUILT,
+ "github_release_assets",
+ lambda repo, tag: {asset_name: "https://example.test/" + asset_name},
+ )
+
+ def fake_download_file(url, destination):
+ raw_calls.append(url)
+
+ def fake_download_file_verified(url, destination, *, expected_sha256, label):
+ verified_calls.append(url)
+
+ monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "download_file", fake_download_file)
+ monkeypatch.setattr(
+ INSTALL_LLAMA_PREBUILT, "download_file_verified", fake_download_file_verified
+ )
+
+ ensure_diffusion_visual_server(
+ tmp_path / "install",
+ linux_host(),
+ "b9334",
+ ApprovedReleaseChecksums(
+ repo = "unslothai/llama.cpp",
+ release_tag = "b9334",
+ upstream_tag = "b9334",
+ artifacts = {},
+ ),
+ )
+
+ target = tmp_path / "install" / "build" / "bin" / "llama-diffusion-gemma-visual-server"
+ assert not target.exists()
+ assert raw_calls == []
+ assert verified_calls == []
From ab6c9ecfee545869d56cc6eddd1babc3f7f36fba Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Wed, 24 Jun 2026 11:37:08 -0300
Subject: [PATCH 55/56] Studio: honor `stream=false` on the GGUF agentic tool
path (#6570) (#6618)
* Studio: honor stream=false on the GGUF agentic tool path (#6570)
* Studio: dedup the #6570 non-streaming tool tests and cover cached_tokens
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: cover the cached_tokens metadata fix and clarify the drain comment (#6570)
* Studio: align the GGUF tool drain naming and tighten its comment (#6570)
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
studio/backend/core/inference/llama_cpp.py | 15 +-
studio/backend/routes/inference.py | 126 ++++++++++++-
.../tests/test_gguf_tool_non_streaming.py | 172 ++++++++++++++++++
.../backend/tests/test_llama_cpp_tool_loop.py | 77 ++++++++
.../tests/test_openai_tool_passthrough.py | 2 +
5 files changed, 378 insertions(+), 14 deletions(-)
create mode 100644 studio/backend/tests/test_gguf_tool_non_streaming.py
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 9c22db4fec..152a3f19b2 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -7881,13 +7881,18 @@ class LlamaCppBackend:
_mt["predicted_per_second"] = _mt["predicted_n"] / (
_mt["predicted_ms"] / 1000.0
)
+ _usage = {
+ "prompt_tokens": _fp,
+ "completion_tokens": _tc,
+ "total_tokens": _fp + _tc,
+ }
+ # Preserve KV-cache hit details (cached_tokens) so the tool path
+ # reports them like the standard non-tool path does, not always 0.
+ if _fu.get("prompt_tokens_details"):
+ _usage["prompt_tokens_details"] = _fu["prompt_tokens_details"]
return {
"type": "metadata",
- "usage": {
- "prompt_tokens": _fp,
- "completion_tokens": _tc,
- "total_tokens": _fp + _tc,
- },
+ "usage": _usage,
"timings": _mt,
"finish_reason": finish_reason,
}
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 2e2c38933e..8b0981cd2b 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -5425,15 +5425,123 @@ async def openai_chat_completions(
pass
_tracker.__exit__(None, None, None)
- return _SameTaskStreamingResponse(
- gguf_tool_stream(),
- media_type = "text/event-stream",
- headers = {
- "Cache-Control": "no-cache",
- "Connection": "close",
- "X-Accel-Buffering": "no",
- },
- )
+ if payload.stream:
+ return _SameTaskStreamingResponse(
+ gguf_tool_stream(),
+ media_type = "text/event-stream",
+ headers = {
+ "Cache-Control": "no-cache",
+ "Connection": "close",
+ "X-Accel-Buffering": "no",
+ },
+ )
+
+ # Non-streaming JSON: drain the agentic generator into one
+ # ChatCompletion, like the standard GGUF `else` branch. stream:false
+ # with tools enabled used to return an SSE body, breaking
+ # non-streaming clients; `unsloth studio run --model` forces tools on
+ # process-wide, so plain requests reach this path (#6570).
+ def _drain_gguf_tool_loop():
+ full_text = ""
+ usage = None
+ finish = None
+ gen = gguf_generate_with_tools()
+ try:
+ for event in gen:
+ if cancel_event.is_set():
+ break
+ if event.get("type") == "metadata":
+ usage = event.get("usage")
+ finish = event.get("finish_reason")
+ elif event.get("type") == "content":
+ # Content is cumulative within a turn and resets
+ # between turns, so the last event holds the final
+ # turn's text. As in the safetensors drain, a visible
+ # preamble emitted before a tool call (its own earlier
+ # turn) isn't carried -- only the final turn is.
+ full_text = _strip_tool_xml_for_display(
+ event.get("text", ""),
+ auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
+ )
+ return full_text, usage, finish
+ finally:
+ # Close the generator on early break/cancel so the underlying
+ # llama-server stream socket is released, like the SSE path.
+ try:
+ gen.close()
+ except (RuntimeError, ValueError):
+ pass
+
+ try:
+ full_text, completion_usage, completion_finish = await asyncio.to_thread(
+ _drain_gguf_tool_loop
+ )
+ reasoning_text, visible_text = _extract_responses_reasoning(
+ full_text,
+ parse_think_markers = _responses_should_parse_think_markers(
+ payload, llama_backend
+ ),
+ )
+ message_kwargs = {"content": visible_text}
+ if reasoning_text:
+ message_kwargs["reasoning_content"] = reasoning_text
+ _usage = completion_usage or {}
+ _prompt_tokens = _usage.get("prompt_tokens") or 0
+ _completion_tokens = _usage.get("completion_tokens") or 0
+ response = ChatCompletion(
+ id = completion_id,
+ created = created,
+ model = model_name,
+ choices = [
+ CompletionChoice(
+ message = CompletionMessage(**message_kwargs),
+ finish_reason = _clamp_finish_reason(completion_finish),
+ )
+ ],
+ usage = CompletionUsage(
+ prompt_tokens = _prompt_tokens,
+ completion_tokens = _completion_tokens,
+ total_tokens = _prompt_tokens + _completion_tokens,
+ prompt_tokens_details = _prompt_tokens_details(
+ _usage.get("prompt_tokens_details")
+ ),
+ ),
+ )
+ api_monitor.set_reply(monitor_id, visible_text)
+ _monitor_usage(
+ monitor_id,
+ {
+ "prompt_tokens": _prompt_tokens,
+ "completion_tokens": _completion_tokens,
+ "total_tokens": _prompt_tokens + _completion_tokens,
+ },
+ _monitor_context_length(),
+ )
+ api_monitor.finish(
+ monitor_id, "cancelled" if cancel_event.is_set() else "completed"
+ )
+ return _model_json_response(response)
+ except Exception as e:
+ logger.error(f"Error during GGUF tool completion: {e}", exc_info = True)
+ api_monitor.fail(monitor_id, _friendly_error(e))
+ # Recover if an MTP+tensor crash killed the server.
+ get_llama_cpp_backend()._maybe_recover_from_mtp_crash(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)
+ if _cls is not None:
+ raise HTTPException(
+ status_code = 400,
+ detail = openai_error_body(
+ _friendly_error(e),
+ status = 400,
+ code = "context_length_exceeded" if _cls else None,
+ param = "messages",
+ ),
+ )
+ raise HTTPException(status_code = 500, detail = safe_error_detail(e))
+ finally:
+ _tracker.__exit__(None, None, None)
# ── Standard GGUF path (no tools) ─────────────────────
diff --git a/studio/backend/tests/test_gguf_tool_non_streaming.py b/studio/backend/tests/test_gguf_tool_non_streaming.py
new file mode 100644
index 0000000000..d9044824cb
--- /dev/null
+++ b/studio/backend/tests/test_gguf_tool_non_streaming.py
@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Regression tests for `stream:false` on the GGUF agentic tool path (#6570).
+
+When server-side tools are enabled (e.g. `unsloth studio run --model ...`,
+which forces the tool policy on process-wide), a plain chat request used to be
+routed into the tool loop, which returned an SSE body *regardless* of
+`stream:false` -- breaking non-streaming clients and health checks like
+LiteLLM. These tests drive the real route with a fake tool-capable backend and
+assert the non-streaming path now returns a single JSON `chat.completion`,
+while `stream:true` still streams.
+"""
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from auth.authentication import get_current_subject
+import routes.inference as inference_route
+
+
+class _ToolGgufBackend:
+ is_loaded = True
+ model_identifier = "test/model.gguf"
+ _is_audio = False
+ is_vision = False
+ supports_tools = True
+
+ def generate_chat_completion_with_tools(self, **kwargs):
+ # The agentic loop runs one tool, then the model answers. Event shapes
+ # mirror the real GGUF loop (tool_start/tool_end/content/metadata).
+ yield {
+ "type": "tool_start",
+ "tool_name": "python",
+ "tool_call_id": "call_1",
+ "arguments": {"code": "print(6 * 7)"},
+ }
+ yield {
+ "type": "tool_end",
+ "tool_name": "python",
+ "tool_call_id": "call_1",
+ "result": "42\n",
+ }
+ yield {"type": "content", "text": "The answer is 42."}
+ yield {
+ "type": "metadata",
+ "usage": {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16},
+ "timings": {"prompt_n": 11, "predicted_n": 5},
+ "finish_reason": "stop",
+ }
+
+
+def _client(monkeypatch, backend = None):
+ monkeypatch.setattr(
+ inference_route, "get_llama_cpp_backend", lambda: backend or _ToolGgufBackend()
+ )
+ # Tools forced on -- the same effect as the CLI `run --model` tool policy.
+ monkeypatch.setattr(inference_route, "_effective_enable_tools", lambda payload: True)
+
+ async def _fake_select(payload, **_kwargs):
+ return [{"type": "function", "function": {"name": "python"}}]
+
+ monkeypatch.setattr(inference_route, "_select_request_tools", _fake_select)
+
+ app = FastAPI()
+ app.include_router(inference_route.router)
+ app.dependency_overrides[get_current_subject] = lambda: "test-user"
+ return TestClient(app)
+
+
+def _payload(stream: bool):
+ return {
+ "messages": [{"role": "user", "content": "What is 6 * 7? Use python."}],
+ "stream": stream,
+ "enable_tools": True,
+ }
+
+
+def test_non_streaming_tool_call_returns_single_json(monkeypatch):
+ response = _client(monkeypatch).post("/chat/completions", json = _payload(stream = False))
+
+ assert response.status_code == 200
+ # The bug returned text/event-stream here; it must be a single JSON object.
+ assert response.headers["content-type"].startswith("application/json")
+
+ body = response.json()
+ assert body["object"] == "chat.completion"
+ choice = body["choices"][0]
+ assert choice["message"]["content"] == "The answer is 42."
+ assert choice["finish_reason"] == "stop"
+ assert body["usage"]["prompt_tokens"] == 11
+ assert body["usage"]["completion_tokens"] == 5
+ assert body["usage"]["total_tokens"] == 16
+
+
+def test_streaming_tool_call_still_streams(monkeypatch):
+ # The parallel path is untouched: stream:true keeps returning SSE.
+ response = _client(monkeypatch).post("/chat/completions", json = _payload(stream = True))
+
+ assert response.status_code == 200
+ assert response.headers["content-type"].startswith("text/event-stream")
+ assert "The answer is 42." in response.text
+ assert "data: [DONE]" in response.text
+
+
+class _EventsBackend(_ToolGgufBackend):
+ """Tool backend that yields a caller-supplied event list."""
+
+ def __init__(self, events):
+ self._events = events
+
+ def generate_chat_completion_with_tools(self, **kwargs):
+ yield from self._events
+
+
+def test_non_streaming_missing_usage_defaults_to_zero(monkeypatch):
+ # No metadata event at all: usage zero-defaults and finish_reason falls back.
+ events = [{"type": "content", "text": "hi"}]
+ response = _client(monkeypatch, _EventsBackend(events)).post(
+ "/chat/completions", json = _payload(stream = False)
+ )
+
+ assert response.status_code == 200
+ body = response.json()
+ assert body["choices"][0]["message"]["content"] == "hi"
+ assert body["choices"][0]["finish_reason"] == "stop"
+ assert body["usage"]["prompt_tokens"] == 0
+ assert body["usage"]["completion_tokens"] == 0
+ assert body["usage"]["total_tokens"] == 0
+
+
+def test_non_streaming_preserves_length_finish_reason(monkeypatch):
+ events = [
+ {"type": "content", "text": "truncated"},
+ {
+ "type": "metadata",
+ "usage": {"prompt_tokens": 3, "completion_tokens": 9},
+ "finish_reason": "length",
+ },
+ ]
+ response = _client(monkeypatch, _EventsBackend(events)).post(
+ "/chat/completions", json = _payload(stream = False)
+ )
+
+ assert response.status_code == 200
+ body = response.json()
+ assert body["choices"][0]["finish_reason"] == "length"
+ # total_tokens is derived when the server omits it.
+ assert body["usage"]["total_tokens"] == 12
+
+
+def test_non_streaming_preserves_cached_tokens(monkeypatch):
+ # KV-cache hit details from the metadata event must survive into the body
+ # (the tool path used to drop them and always report cached_tokens=0).
+ events = [
+ {"type": "content", "text": "hi"},
+ {
+ "type": "metadata",
+ "usage": {
+ "prompt_tokens": 20,
+ "completion_tokens": 4,
+ "prompt_tokens_details": {"cached_tokens": 16},
+ },
+ "finish_reason": "stop",
+ },
+ ]
+ response = _client(monkeypatch, _EventsBackend(events)).post(
+ "/chat/completions", json = _payload(stream = False)
+ )
+
+ assert response.status_code == 200
+ assert response.json()["usage"]["prompt_tokens_details"]["cached_tokens"] == 16
diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py
index 56e028bd5a..05d2a0b80a 100644
--- a/studio/backend/tests/test_llama_cpp_tool_loop.py
+++ b/studio/backend/tests/test_llama_cpp_tool_loop.py
@@ -1736,3 +1736,80 @@ def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch):
assert provisional == []
# The real call still executes despite the missing id.
assert calls == [("python", {"code": big_code})]
+
+
+def _usage_done(usage: dict, finish_reason: str = "stop") -> str:
+ """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the
+ real server reports it on the final chunk of a completion."""
+ return (
+ "data: "
+ + json.dumps(
+ {
+ "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
+ "usage": usage,
+ }
+ )
+ + "\n"
+ )
+
+
+def test_metadata_event_preserves_prompt_tokens_details(monkeypatch):
+ """The tool loop's metadata event must carry llama-server's
+ ``prompt_tokens_details`` (KV-cache hits) through ``_build_metadata_event``,
+ so the route reports real ``cached_tokens`` instead of always 0 (#6570).
+
+ This drives the *real* generator; the route-level test feeds a pre-built
+ metadata event and so never exercises this code.
+ """
+ stream = [
+ _sse({"content": "The answer is 42."}),
+ _usage_done(
+ {
+ "prompt_tokens": 20,
+ "completion_tokens": 4,
+ "prompt_tokens_details": {"cached_tokens": 16},
+ }
+ ),
+ _done(),
+ ]
+ payloads: list[dict] = []
+ backend = _make_backend(monkeypatch, [stream], payloads)
+
+ events = list(
+ backend.generate_chat_completion_with_tools(
+ messages = [{"role": "user", "content": "hi"}],
+ tools = [],
+ max_tool_iterations = 1,
+ )
+ )
+
+ metadata = [e for e in events if e.get("type") == "metadata"]
+ assert metadata, "expected a metadata event"
+ usage = metadata[-1]["usage"]
+ assert usage["prompt_tokens_details"] == {"cached_tokens": 16}
+ assert usage["prompt_tokens"] == 20
+ assert usage["completion_tokens"] == 4
+
+
+def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch):
+ """No KV-cache block from the server -> the key isn't fabricated, so the
+ route falls back to its 0-default instead of reading a bogus value."""
+ stream = [
+ _sse({"content": "hi"}),
+ _usage_done({"prompt_tokens": 5, "completion_tokens": 2}),
+ _done(),
+ ]
+ payloads: list[dict] = []
+ backend = _make_backend(monkeypatch, [stream], payloads)
+
+ events = list(
+ backend.generate_chat_completion_with_tools(
+ messages = [{"role": "user", "content": "hi"}],
+ tools = [],
+ max_tool_iterations = 1,
+ )
+ )
+
+ metadata = [e for e in events if e.get("type") == "metadata"]
+ assert metadata, "expected a metadata event"
+ assert "prompt_tokens_details" not in metadata[-1]["usage"]
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index aaef9e4dcc..aa36c6fed4 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -1349,6 +1349,7 @@ class TestGgufVisionToolRouting:
model = "default",
enable_tools = True,
enabled_tools = ["web_search"],
+ stream = True,
messages = [
{
"role": "user",
@@ -1408,6 +1409,7 @@ class TestGgufVisionToolRouting:
enable_tools = True,
enabled_tools = ["web_search"],
parallel_tool_calls = False,
+ stream = True,
messages = [{"role": "user", "content": "search once"}],
)
From a3954edd15e4a03b584d60940173b99c17f45922 Mon Sep 17 00:00:00 2001
From: Long Yixing
Date: Wed, 24 Jun 2026 22:39:20 +0800
Subject: [PATCH 56/56] Fix Studio GGUF variant expansion crash (#6636)
* fix: handle empty GGUF variants
* fix: gate local GGUF expansion
* fix: normalize GGUF variant payload
---------
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
.../assistant-ui/model-selector/pickers.tsx | 83 +++++++++++++++----
1 file changed, 68 insertions(+), 15 deletions(-)
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 90c1c04106..386f233c7b 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -570,6 +570,52 @@ function ModelRow({
// ── GGUF Variant Expander ────────────────────────────────────
+function isValidGgufVariant(variant: unknown): variant is GgufVariantDetail {
+ if (!variant || typeof variant !== "object") return false;
+ const candidate = variant as Partial;
+ return (
+ typeof candidate.filename === "string" &&
+ candidate.filename.length > 0 &&
+ typeof candidate.quant === "string" &&
+ candidate.quant.length > 0 &&
+ typeof candidate.size_bytes === "number" &&
+ Number.isFinite(candidate.size_bytes) &&
+ candidate.size_bytes >= 0 &&
+ (candidate.downloaded === undefined ||
+ typeof candidate.downloaded === "boolean")
+ );
+}
+
+function normalizeGgufVariantsResponse(res: {
+ variants?: unknown;
+ default_variant?: unknown;
+ has_vision?: unknown;
+ context_length?: unknown;
+} | null | undefined): {
+ variants: GgufVariantDetail[];
+ defaultVariant: string | null;
+ hasVision: boolean;
+ contextLength: number | null;
+} {
+ const contextLength = res?.context_length;
+ return {
+ variants: (Array.isArray(res?.variants) ? res.variants : []).filter(
+ isValidGgufVariant,
+ ),
+ defaultVariant:
+ typeof res?.default_variant === "string" && res.default_variant.length > 0
+ ? res.default_variant
+ : null,
+ hasVision: res?.has_vision === true,
+ contextLength:
+ typeof contextLength === "number" &&
+ Number.isFinite(contextLength) &&
+ contextLength >= 0
+ ? contextLength
+ : null,
+ };
+}
+
function GgufVariantExpander({
repoId,
onSelect,
@@ -622,11 +668,12 @@ function GgufVariantExpander({
listGgufVariants(repoId)
.then((res) => {
if (canceled) return;
- setVariants(res.variants);
- setDefaultVariant(res.default_variant);
- setHasVision(res.has_vision);
- onHasVision?.(res.has_vision);
- setNativeContext(res.context_length ?? null);
+ const normalized = normalizeGgufVariantsResponse(res);
+ setVariants(normalized.variants);
+ setDefaultVariant(normalized.defaultVariant);
+ setHasVision(normalized.hasVision);
+ onHasVision?.(normalized.hasVision);
+ setNativeContext(normalized.contextLength);
})
.catch((err) => {
if (canceled) return;
@@ -694,19 +741,25 @@ function GgufVariantExpander({
// If the recommended variant is OOM, pick the largest fitting one;
// if all are OOM, recommend the smallest.
const effectiveRecommended = useMemo(() => {
- if (!variants || totalBudgetGb <= 0) return defaultVariant;
+ if (!variants || variants.length === 0 || totalBudgetGb <= 0) {
+ return defaultVariant;
+ }
const defaultV = variants.find((v) => v.quant === defaultVariant);
if (defaultV && getGgufFit(defaultV.size_bytes) !== "oom")
return defaultVariant;
// Largest non-OOM variant (best quality that fits)
- const fitting = variants.filter((v) => getGgufFit(v.size_bytes) !== "oom");
+ const fitting = variants.filter(
+ (v) => getGgufFit(v.size_bytes) !== "oom",
+ );
if (fitting.length > 0) {
fitting.sort((a, b) => b.size_bytes - a.size_bytes);
return fitting[0].quant;
}
// All OOM -- recommend smallest (most likely to partially run)
- const sorted = [...variants].sort((a, b) => a.size_bytes - b.size_bytes);
- return sorted[0].quant;
+ const sorted = [...variants].sort(
+ (a, b) => a.size_bytes - b.size_bytes,
+ );
+ return sorted[0]?.quant ?? defaultVariant;
}, [variants, defaultVariant, totalBudgetGb, getGgufFit]);
const sortedVariants = useMemo(() => {
@@ -2901,7 +2954,7 @@ export function HubModelPicker({
}
}}
onArrowDownIntoChildren={
- isGgufExpanded(m.id)
+ isGguf && !isDirectGguf && isGgufExpanded(m.id)
? () => {
const focused =
focusFirstChildOption(optionKey);
@@ -2911,7 +2964,7 @@ export function HubModelPicker({
}
vramStatus={null}
/>
- {isGgufExpanded(m.id) && (
+ {isGguf && !isDirectGguf && isGgufExpanded(m.id) && (
{
const focused =
focusFirstChildOption(optionKey);
@@ -2998,7 +3051,7 @@ export function HubModelPicker({
}
vramStatus={null}
/>
- {!isGgufFile && isGgufExpanded(m.id) && (
+ {isGguf && !isGgufFile && isGgufExpanded(m.id) && (
focusFirstChildOption(optionKey)
: undefined
}
vramStatus={null}
/>
- {!isGgufFile && isGgufExpanded(m.id) && (
+ {isGguf && !isGgufFile && isGgufExpanded(m.id) && (