Studio: harden OpenAI-compatible GGUF streaming (#6950)

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>

studio/backend/core/inference/llama_cpp.py

@@ -100,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError):
     Subclasses RuntimeError so existing handlers still catch it."""
 
 
+class _LlamaStreamCancelled(Exception):
+    """Internal signal for an expected client/request cancellation."""
+
+
 # Shared so the from_identifier preflight and the load-time raise stay in sync.
 LLAMA_SERVER_NOT_FOUND_DETAIL = (
     "This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
@@ -8616,7 +8620,7 @@ class LlamaCppBackend:
     ):
         """Open one streaming POST and let cancel interrupt prefill or reads."""
         if cancel_event is not None and cancel_event.is_set():
-            raise GeneratorExit
+            raise _LlamaStreamCancelled
 
         _cancel_closed = threading.Event()
         _response_ref: list = [None]
@@ -8661,13 +8665,13 @@ class LlamaCppBackend:
             ) as response:
                 _response_ref[0] = response
                 if cancel_event is not None and cancel_event.is_set():
-                    raise GeneratorExit
+                    raise _LlamaStreamCancelled
                 yield response
                 return
         except (httpx.RequestError, RuntimeError):
             # Response was closed by the cancel watcher
             if cancel_event is not None and cancel_event.is_set():
-                raise GeneratorExit
+                raise _LlamaStreamCancelled
             raise
         finally:
             _cancel_closed.set()
@@ -8870,6 +8874,8 @@ class LlamaCppBackend:
                         "finish_reason": _metadata_finish_reason,
                     }
 
+        except _LlamaStreamCancelled:
+            return
         except httpx.ConnectError as e:
             # Server already down. If this was an MTP+tensor crash, recover by
             # reloading without MTP (scheduled in the background) and fail this
@@ -9994,6 +10000,8 @@ class LlamaCppBackend:
                         break
                 continue
 
+            except _LlamaStreamCancelled:
+                return
             except httpx.ConnectError:
                 # Mark unresolved provisional cards as failed before raising.
                 for _pid, _pname in provisional_started_tool_calls.items():
@@ -10176,6 +10184,8 @@ class LlamaCppBackend:
                 if _meta is not None:
                     yield _meta
 
+        except _LlamaStreamCancelled:
+            return
         except httpx.ConnectError:
             raise RuntimeError("Lost connection to llama-server")
         except Exception as e:

studio/backend/tests/test_inference_dispatcher_resilience.py

@@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
             continue
         calls.append(node)
 
-    assert len(calls) == 4
+    assert len(calls) == 5
     for call in calls:
         assert any(
             kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False

studio/backend/tests/test_llama_cpp_stream_cancel.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+import contextlib
+import os
+import sys
+import threading
+
+import httpx
+import pytest
+
+_backend = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, _backend)
+
+from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled
+
+
+def _backend_stub() -> LlamaCppBackend:
+    backend = LlamaCppBackend.__new__(LlamaCppBackend)
+    backend._process = object()
+    backend._healthy = True
+    backend._port = 48848
+    backend._effective_context_length = 4096
+    backend._supports_reasoning = False
+    backend._reasoning_always_on = False
+    backend._reasoning_style = "enable_thinking"
+    backend._supports_preserve_thinking = False
+    return backend
+
+
+def test_stream_cancel_uses_internal_exception_not_generator_exit():
+    class FakeResponse:
+        status_code = 200
+
+        def close(self):
+            pass
+
+    class FakeStream:
+        def __enter__(self):
+            return FakeResponse()
+
+        def __exit__(self, *_args):
+            return False
+
+    class FakeClient:
+        def stream(self, *_args, **_kwargs):
+            return FakeStream()
+
+    cancel_event = threading.Event()
+
+    with pytest.raises(Exception) as exc_info:
+        with LlamaCppBackend._stream_with_retry(
+            FakeClient(),
+            "http://llama.test/v1/chat/completions",
+            {},
+            cancel_event,
+        ):
+            cancel_event.set()
+            raise httpx.ReadError("client closed")
+
+    assert exc_info.type is _LlamaStreamCancelled
+    assert not issubclass(exc_info.type, GeneratorExit)
+
+
+def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch):
+    backend = _backend_stub()
+
+    @contextlib.contextmanager
+    def fake_open_stream(*_args, **_kwargs):
+        raise _LlamaStreamCancelled
+
+    monkeypatch.setattr(backend, "_open_stream", fake_open_stream)
+
+    chunks = list(
+        backend.generate_chat_completion(
+            [{"role": "user", "content": "hi"}],
+            cancel_event = threading.Event(),
+        )
+    )
+
+    assert chunks == []

studio/backend/tests/test_llama_route_timeouts.py

@@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel():
 
     asyncio.run(_run(False))
     asyncio.run(_run(True))
+
+
+def test_stream_stall_timeout_callable_re_resolved_each_read():
+    # The OpenAI passthrough passes a callable so the stall bound can switch to
+    # the short post-terminal grace mid-stream; it must be re-resolved per read,
+    # not captured once at generator start.
+    async def _run():
+        response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
+        values = iter([100.0, 2.0])
+        seen = []
+
+        class _Request:
+            async def is_disconnected(self):
+                return False
+
+        class _Items:
+            def __init__(self):
+                self.count = 0
+
+            async def __anext__(self):
+                self.count += 1
+                if self.count > 3:
+                    raise StopAsyncIteration
+                return "data: {}"
+
+        async for _ in inf_mod._aiter_llama_stream_items(
+            _Items(),
+            cancel_event = threading.Event(),
+            request = _Request(),
+            response = response,
+            first_token_deadline = time.monotonic() + 1,
+            post_first_item_read_timeout_s = lambda: next(values, 5.0),
+        ):
+            seen.append(response.request.extensions["timeout"].get("read"))
+
+        assert len(seen) == 3
+        # The callable is resolved right after the first item (arming the
+        # post-first window) and again before each later read, consuming
+        # successive values.
+        assert seen[0] == 100.0
+        assert 1.0 <= seen[1] <= 2.0
+        assert 4.0 <= seen[2] <= 5.0
+
+    asyncio.run(_run())
+
+
+def test_stream_stall_timeout_disabled_clears_read_timeout():
+    # UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so
+    # the callable returns None. Once a chunk has arrived the leftover
+    # first-token read timeout must be cleared, else a long post-first-chunk gap
+    # trips a stale deadline the operator asked to turn off.
+    async def _run():
+        response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
+        seen = []
+
+        class _Request:
+            async def is_disconnected(self):
+                return False
+
+        class _Items:
+            def __init__(self):
+                self.count = 0
+
+            async def __anext__(self):
+                self.count += 1
+                if self.count > 2:
+                    raise StopAsyncIteration
+                return "data: {}"
+
+        async for _ in inf_mod._aiter_llama_stream_items(
+            _Items(),
+            cancel_event = threading.Event(),
+            request = _Request(),
+            response = response,
+            first_token_deadline = time.monotonic() + 5,
+            post_first_item_read_timeout_s = lambda: None,
+        ):
+            seen.append(response.request.extensions["timeout"].get("read"))
+
+        # The first-token path armed a finite read timeout; after the first chunk
+        # with the guard disabled, it is cleared to None on every subsequent read.
+        assert seen == [None, None], seen
+
+    asyncio.run(_run())

studio/backend/tests/test_passthrough_healing.py

@@ -515,6 +515,7 @@ class ScriptedClient:
         _url,
         json = None,
         timeout = None,
+        headers = None,
     ):
         self.posts.append(json)
         return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])