From ab6c9ecfee545869d56cc6eddd1babc3f7f36fba Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Wed, 24 Jun 2026 11:37:08 -0300
Subject: [PATCH] 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 9c22db4fe..152a3f19b 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 2e2c38933..8b0981cd2 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 000000000..d9044824c
--- /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 56e028bd5..05d2a0b80 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 aaef9e4dc..aa36c6fed 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"}],
         )