[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

studio/backend/core/inference/llama_cpp.py

@@ -57,6 +57,7 @@ class LlamaCppBackend:
         self._stdout_lines: list[str] = []
         self._stdout_thread: Optional[threading.Thread] = None
         self._cancel_event = threading.Event()
+        self._api_key: Optional[str] = None
 
         self._kill_orphaned_servers()
         atexit.register(self._cleanup)
@@ -938,6 +939,17 @@ class LlamaCppBackend:
                     cmd.extend(["--mmproj", mmproj_path])
                     logger.info(f"Using mmproj for vision: {mmproj_path}")
 
+            # Option C: add --api-key for direct client access when enabled
+            import os as _os
+            import secrets as _secrets
+
+            if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1":
+                self._api_key = _secrets.token_urlsafe(32)
+                cmd.extend(["--api-key", self._api_key])
+                logger.info("llama-server started with --api-key for direct streaming")
+            else:
+                self._api_key = None
+
             logger.info(f"Starting llama-server: {' '.join(cmd)}")
 
             # Set library paths so llama-server can find its shared libs and CUDA DLLs
@@ -1407,6 +1419,7 @@ class LlamaCppBackend:
         url: str,
         payload: dict,
         cancel_event: Optional[threading.Event] = None,
+        headers: Optional[dict] = None,
     ):
         """Open an httpx streaming POST with cancel support.
 
@@ -1473,7 +1486,11 @@ class LlamaCppBackend:
                 pool = 10,
             )
             with client.stream(
-                "POST", url, json = payload, timeout = prefill_timeout
+                "POST",
+                url,
+                json = payload,
+                timeout = prefill_timeout,
+                headers = headers,
             ) as response:
                 _response_ref[0] = response
                 if cancel_event is not None and cancel_event.is_set():
@@ -1547,9 +1564,16 @@ class LlamaCppBackend:
             # can finish.  Cancel during streaming is handled by the
             # watcher thread (closes the response on cancel_event).
             stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
+            _auth_headers = (
+                {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
+            )
             with httpx.Client(timeout = stream_timeout) as client:
                 with self._stream_with_retry(
-                    client, url, payload, cancel_event
+                    client,
+                    url,
+                    payload,
+                    cancel_event,
+                    headers = _auth_headers,
                 ) as response:
                     if response.status_code != 200:
                         error_body = response.read().decode()
@@ -1706,8 +1730,13 @@ class LlamaCppBackend:
                 payload["stop"] = stop
 
             try:
+                _auth_headers = (
+                    {"Authorization": f"Bearer {self._api_key}"}
+                    if self._api_key
+                    else None
+                )
                 with httpx.Client(timeout = None) as client:
-                    resp = client.post(url, json = payload)
+                    resp = client.post(url, json = payload, headers = _auth_headers)
                     if resp.status_code != 200:
                         raise RuntimeError(
                             f"llama-server returned {resp.status_code}: {resp.text}"
@@ -1950,9 +1979,16 @@ class LlamaCppBackend:
 
         try:
             stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
+            _auth_headers = (
+                {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
+            )
             with httpx.Client(timeout = stream_timeout) as client:
                 with self._stream_with_retry(
-                    client, url, stream_payload, cancel_event
+                    client,
+                    url,
+                    stream_payload,
+                    cancel_event,
+                    headers = _auth_headers,
                 ) as response:
                     if response.status_code != 200:
                         error_body = response.read().decode()
@@ -2078,7 +2114,10 @@ class LlamaCppBackend:
         if not self.is_loaded:
             return None
         try:
-            with httpx.Client(timeout = 10) as client:
+            _auth_headers = (
+                {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
+            )
+            with httpx.Client(timeout = 10, headers = _auth_headers) as client:
 
                 def _detok(tid: int) -> str:
                     r = client.post(
@@ -2196,7 +2235,12 @@ class LlamaCppBackend:
         if need_ids:
             payload["n_probs"] = 1
 
-        with httpx.Client(timeout = httpx.Timeout(300, connect = 10)) as client:
+        _auth_headers = (
+            {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
+        )
+        with httpx.Client(
+            timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers
+        ) as client:
             resp = client.post(f"{self.base_url}/completion", json = payload)
             if resp.status_code != 200:
                 raise RuntimeError(

studio/backend/main.py

@@ -112,6 +112,27 @@ async def lifespan(app: FastAPI):
         print("=" * 60 + "\n")
     else:
         app.state.bootstrap_password = storage.get_bootstrap_password()
+
+    # Start dedicated streaming server in a daemon thread (Option B).
+    # Accepts both UNSLOTH_FAST_SSE=1 and UNSLOTH_STREAM_SERVER=1.
+    if (
+        os.getenv("UNSLOTH_FAST_SSE", "0") == "1"
+        or os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1"
+    ):
+        import threading as _threading
+
+        from streaming_server import start_streaming_server, find_free_port
+
+        stream_port = find_free_port()
+        app.state.stream_port = stream_port
+        _stream_thread = _threading.Thread(
+            target = start_streaming_server,
+            args = (stream_port,),
+            daemon = True,
+        )
+        _stream_thread.start()
+        print(f"[streaming_server] Started on 127.0.0.1:{stream_port}")
+
     yield
     # Cleanup
     _hw_module.DEVICE = None
@@ -373,3 +394,9 @@ def setup_frontend(app: FastAPI, build_path: Path):
         )
 
     return True
+
+
+# Note: Option A (RawSSEInterceptor / asgi_fast_path.py) has been removed.
+# Benchmarking proved it ineffective (~172 TPS = baseline) since it shares the
+# same asyncio event loop. UNSLOTH_FAST_SSE=1 now starts the dedicated streaming
+# server (Option B) which runs in a separate thread with its own event loop.

studio/backend/routes/inference.py

@@ -1563,6 +1563,138 @@ async def openai_chat_completions(
 # =====================================================================
 
 
+# =====================================================================
+# Option B: Stream URL endpoint (one-time token for streaming server)
+# =====================================================================
+
+
+@router.get("/stream-url")
+async def get_stream_url(
+    request: Request,
+    current_subject: str = Depends(get_current_subject),
+):
+    """
+    Issue a one-time streaming token and return the URL for the dedicated
+    streaming server (Option B).
+
+    Requires UNSLOTH_FAST_SSE=1 or UNSLOTH_STREAM_SERVER=1.
+
+    Returns ``{"supported": false}`` when the fast path cannot serve the
+    current model (no GGUF loaded, audio model active, or streaming server
+    not enabled). The frontend uses this to decide whether to use the fast
+    path or fall back to the baseline ``/v1/chat/completions``.
+
+    When supported, returns the stream URL (token NOT in query string --
+    the client must pass it via the ``X-Stream-Token`` header).
+    """
+    import os
+
+    fast_sse = os.getenv("UNSLOTH_FAST_SSE", "0") == "1"
+    stream_server = os.getenv("UNSLOTH_STREAM_SERVER", "0") == "1"
+    if not fast_sse and not stream_server:
+        return {"supported": False}
+
+    stream_port = getattr(request.app.state, "stream_port", None)
+    if stream_port is None:
+        return {"supported": False}
+
+    # Check if current model is compatible with fast path
+    llama_backend = get_llama_cpp_backend()
+    if not llama_backend.is_loaded:
+        return {"supported": False}
+    if getattr(llama_backend, "_is_audio", False):
+        return {"supported": False}
+
+    from stream_token_store import create_stream_token
+
+    token = create_stream_token(current_subject)
+    return {
+        "supported": True,
+        "stream_url": f"http://127.0.0.1:{stream_port}/stream",
+        "token": token,
+        "port": stream_port,
+        "ttl_seconds": 10,
+    }
+
+
+# =====================================================================
+# Internal: Stream token validation for streaming server subprocess
+# =====================================================================
+
+
+@router.post("/internal/consume-stream-token")
+async def consume_stream_token_endpoint(request: Request):
+    """
+    Validate a one-time stream token and return llama-server connection info.
+
+    Called by the streaming server subprocess to validate tokens without
+    needing shared memory.  Localhost-only (no auth required).
+    """
+    body = await request.json()
+    token = body.get("token")
+    if not token:
+        return JSONResponse({"valid": False}, status_code = 400)
+
+    from stream_token_store import consume_stream_token
+
+    username = consume_stream_token(token)
+    if not username:
+        return JSONResponse({"valid": False}, status_code = 401)
+
+    llama = get_llama_cpp_backend()
+    return JSONResponse(
+        {
+            "valid": True,
+            "username": username,
+            "llama_port": llama._port if llama.is_loaded else None,
+            "llama_api_key": llama._api_key,
+            "model_name": llama.model_identifier or "unknown",
+            "supports_reasoning": llama.supports_reasoning,
+            "supports_tools": llama.supports_tools,
+            "is_vision": llama.is_vision,
+        }
+    )
+
+
+# =====================================================================
+# Option C: Direct stream endpoint (llama-server with --api-key)
+# =====================================================================
+
+
+@router.get("/direct-stream")
+async def get_direct_stream(
+    current_subject: str = Depends(get_current_subject),
+):
+    """
+    Return the internal llama-server URL and API key for direct client streaming
+    (Option C). This bypasses all Studio transformations (thinking tags, image
+    normalization, cumulative-to-delta) but achieves maximum TPS.
+
+    Requires the llama-server to have been started with --api-key.
+    """
+    llama_backend = get_llama_cpp_backend()
+    if not llama_backend.is_loaded:
+        raise HTTPException(status_code = 400, detail = "No GGUF model loaded.")
+
+    api_key = getattr(llama_backend, "_api_key", None)
+    if not api_key:
+        raise HTTPException(
+            status_code = 501,
+            detail = "llama-server was not started with --api-key. Reload the model.",
+        )
+
+    return {
+        "base_url": llama_backend.base_url,
+        "api_key": api_key,
+        "model": llama_backend.model_identifier,
+    }
+
+
+# =====================================================================
+# OpenAI-Compatible Models Listing  (/models -> /v1/models)
+# =====================================================================
+
+
 @router.get("/models")
 async def openai_list_models(
     current_subject: str = Depends(get_current_subject),

studio/backend/stream_token_store.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+Thread-safe one-time token store for Option B (separate streaming server).
+
+Tokens are short-lived (10 seconds) and consumed on first use.
+"""
+
+import threading
+import time
+import uuid
+from typing import Optional
+
+
+class StreamTokenStore:
+    """Thread-safe store for short-lived, one-time-use streaming tokens."""
+
+    def __init__(self, ttl_seconds: float = 10.0) -> None:
+        self._ttl = ttl_seconds
+        self._lock = threading.Lock()
+        # token -> {"username": str, "expires": float}
+        self._tokens: dict[str, dict] = {}
+
+    def create_token(self, username: str) -> str:
+        """Create a new one-time token for the given user. Returns the token string."""
+        token = uuid.uuid4().hex
+        expires = time.monotonic() + self._ttl
+        with self._lock:
+            self._purge_expired()
+            self._tokens[token] = {"username": username, "expires": expires}
+        return token
+
+    def consume_token(self, token: str) -> Optional[str]:
+        """
+        Validate and consume a token. Returns the username if valid, None otherwise.
+        The token is deleted after consumption (one-time use).
+        """
+        with self._lock:
+            self._purge_expired()
+            entry = self._tokens.pop(token, None)
+        if entry is None:
+            return None
+        if time.monotonic() > entry["expires"]:
+            return None
+        return entry["username"]
+
+    def _purge_expired(self) -> None:
+        """Remove expired tokens. Must be called while holding _lock."""
+        now = time.monotonic()
+        expired = [k for k, v in self._tokens.items() if now > v["expires"]]
+        for k in expired:
+            del self._tokens[k]
+
+
+# Module-level singleton
+_store = StreamTokenStore()
+
+
+def create_stream_token(username: str) -> str:
+    return _store.create_token(username)
+
+
+def consume_stream_token(token: str) -> Optional[str]:
+    return _store.consume_token(token)

studio/backend/streaming_server.py

@@ -0,0 +1,667 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+Dedicated streaming server for fast SSE (Option B).
+
+Runs as a standalone FastAPI app in a separate thread with its own event loop,
+eliminating asyncio contention with the main Studio app.
+
+Authentication: one-time tokens issued by the main app's /stream-url endpoint,
+passed via the X-Stream-Token header (not in the URL to avoid logging leaks).
+
+Supports: streaming, non-streaming, tool calling, vision, thinking mode.
+Full feature parity with baseline /v1/chat/completions.
+
+PERFORMANCE: The streaming hot path uses httpx.AsyncClient to stream directly
+from llama-server, bypassing the sync generate_chat_completion() generator.
+Only sends sampling parameters the client explicitly provides -- notably,
+repeat_penalty defaults to llama-server's own 1.0 instead of being forced
+to 1.1, which avoids a ~24% TPS penalty from repetition scanning.
+"""
+
+import asyncio
+import json
+import re
+import threading
+import time
+import uuid
+from typing import Optional
+
+import httpx
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from stream_token_store import consume_stream_token
+
+
+stream_app = FastAPI(docs_url = None, redoc_url = None, openapi_url = None)
+
+# ── Shared helpers ────────────────────────────────────────────
+
+
+def _friendly_error(e):
+    """Convert raw exception messages to user-readable strings."""
+    msg = str(e)
+    m = re.search(
+        r"request \((\d+) tokens?\) exceeds the available context size \((\d+) tokens?\)",
+        msg,
+    )
+    if m:
+        return (
+            f"Message too long: {m.group(1)} tokens exceeds the {m.group(2)}-token "
+            f"context window. Try increasing the Context Length in Model settings, "
+            f"or shorten the conversation."
+        )
+    if "Lost connection to llama-server" in msg:
+        return "Lost connection to the model server. It may have crashed -- try reloading the model."
+    return "An internal error occurred"
+
+
+def _extract_content_parts(messages):
+    """Parse messages, extracting text and image_b64 from content parts."""
+    gguf_messages = []
+    image_b64 = None
+    for msg in messages:
+        role = msg.get("role", "")
+        content = msg.get("content", "")
+        if isinstance(content, list):
+            text_parts = []
+            for part in content:
+                if isinstance(part, dict):
+                    if part.get("type") == "text":
+                        text_parts.append(part.get("text", ""))
+                    elif part.get("type") == "image_url" and image_b64 is None:
+                        url = part.get("image_url", {}).get("url", "")
+                        if url.startswith("data:") and "," in url:
+                            image_b64 = url.split(",", 1)[1]
+            content = "\n".join(text_parts) if text_parts else ""
+        gguf_messages.append({"role": role, "content": content})
+    return gguf_messages, image_b64
+
+
+def _process_image(image_b64, llama_backend):
+    """Validate and convert image to PNG if needed."""
+    if image_b64 and llama_backend.is_vision:
+        import base64 as _b64
+        from io import BytesIO as _BytesIO
+        from PIL import Image as _Image
+
+        raw = _b64.b64decode(image_b64)
+        img = _Image.open(_BytesIO(raw))
+        if img.mode == "RGBA":
+            img = img.convert("RGB")
+        buf = _BytesIO()
+        img.save(buf, format = "PNG")
+        return _b64.b64encode(buf.getvalue()).decode("ascii")
+    elif image_b64 and not llama_backend.is_vision:
+        raise HTTPException(
+            status_code = 400,
+            detail = "Image provided but current GGUF model does not support vision.",
+        )
+    return image_b64
+
+
+def _build_llama_payload(llama_backend, openai_messages, payload, stream = True):
+    """Build the payload for llama-server /v1/chat/completions.
+
+    Only sends repeat_penalty when the client explicitly provides it.
+    This avoids the ~24% TPS penalty from repetition scanning when
+    the frontend has not set a repetition penalty.
+    """
+    llama_payload = {
+        "messages": openai_messages,
+        "stream": stream,
+        "temperature": payload.get("temperature", 0.6),
+        "top_p": payload.get("top_p", 0.95),
+        "top_k": max(payload.get("top_k", 20), 0),
+        "min_p": payload.get("min_p", 0.0),
+        "presence_penalty": payload.get("presence_penalty", 0.0),
+    }
+    # Only send repeat_penalty when the client explicitly sets repetition_penalty.
+    # llama-server defaults to 1.0; forcing 1.1 costs ~24% TPS.
+    if "repetition_penalty" in payload:
+        llama_payload["repeat_penalty"] = payload["repetition_penalty"]
+    if stream:
+        llama_payload["stream_options"] = {"include_usage": True}
+    if llama_backend.supports_reasoning and payload.get("enable_thinking") is not None:
+        llama_payload["chat_template_kwargs"] = {
+            "enable_thinking": payload["enable_thinking"]
+        }
+    if payload.get("max_tokens") is not None:
+        llama_payload["max_tokens"] = payload["max_tokens"]
+    if payload.get("stop"):
+        llama_payload["stop"] = payload["stop"]
+    return llama_payload
+
+
+# ── CORS preflight ────────────────────────────────────────────
+
+
+@stream_app.options("/stream")
+async def stream_preflight():
+    """Handle CORS preflight for the /stream endpoint."""
+    return JSONResponse(
+        content = {},
+        headers = {
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Methods": "POST, OPTIONS",
+            "Access-Control-Allow-Headers": "Content-Type, X-Stream-Token",
+            "Access-Control-Max-Age": "86400",
+        },
+    )
+
+
+# ── Request validation ────────────────────────────────────────
+
+
+async def _validate_request(request: Request):
+    """Validate token, parse body, get backend. Returns all needed context."""
+    token = request.headers.get("X-Stream-Token")
+    if not token:
+        raise HTTPException(status_code = 401, detail = "Missing X-Stream-Token header")
+    username = consume_stream_token(token)
+    if username is None:
+        raise HTTPException(status_code = 401, detail = "Invalid or expired stream token")
+
+    body_bytes = await request.body()
+    try:
+        payload = json.loads(body_bytes)
+    except (json.JSONDecodeError, UnicodeDecodeError):
+        raise HTTPException(status_code = 400, detail = "Invalid JSON body")
+
+    from routes.inference import get_llama_cpp_backend
+
+    llama_backend = get_llama_cpp_backend()
+    if not llama_backend.is_loaded:
+        raise HTTPException(status_code = 400, detail = "No GGUF model loaded")
+
+    messages = payload.get("messages", [])
+    gguf_messages, image_b64 = _extract_content_parts(messages)
+
+    # Legacy image_base64 fallback
+    if not image_b64:
+        image_b64 = payload.get("image_base64")
+
+    # Image validation and conversion
+    image_b64 = _process_image(image_b64, llama_backend)
+
+    completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+    created = int(time.time())
+    model_name = llama_backend.model_identifier or "unknown"
+
+    return (
+        payload,
+        llama_backend,
+        gguf_messages,
+        image_b64,
+        completion_id,
+        created,
+        model_name,
+    )
+
+
+# ── Path A: Direct async streaming (HOT PATH) ────────────────
+
+
+_SSE_HEADERS = {
+    "Cache-Control": "no-cache",
+    "Connection": "keep-alive",
+    "X-Accel-Buffering": "no",
+    "Access-Control-Allow-Origin": "*",
+}
+
+_SEP = (",", ":")
+
+
+async def _handle_async_stream(
+    request,
+    payload,
+    llama_backend,
+    gguf_messages,
+    image_b64,
+    completion_id,
+    created,
+    model_name,
+):
+    """
+    Stream directly from llama-server using httpx.AsyncClient.
+
+    Bypasses the sync generate_chat_completion() generator and its
+    asyncio.to_thread overhead.  llama-server speaks standard OpenAI SSE
+    with delta tokens natively, so no cumulative-to-delta conversion needed.
+    """
+    openai_messages = llama_backend._build_openai_messages(gguf_messages, image_b64)
+    llama_payload = _build_llama_payload(
+        llama_backend, openai_messages, payload, stream = True
+    )
+
+    port = llama_backend._port
+    api_key = llama_backend._api_key
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+    url = f"http://127.0.0.1:{port}/v1/chat/completions"
+    timeout = httpx.Timeout(connect = 30, read = 120.0, write = 10, pool = 10)
+
+    async def sse_generator():
+        try:
+            # Role chunk
+            role = {
+                "id": completion_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model_name,
+                "choices": [
+                    {"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}
+                ],
+            }
+            yield f"data: {json.dumps(role, separators = _SEP)}\n\n"
+
+            async with httpx.AsyncClient(timeout = timeout) as client:
+                async with client.stream(
+                    "POST", url, json = llama_payload, headers = headers
+                ) as resp:
+                    if resp.status_code != 200:
+                        error_body = await resp.aread()
+                        raise RuntimeError(
+                            f"llama-server returned {resp.status_code}: {error_body.decode()}"
+                        )
+
+                    buffer = ""
+                    in_thinking = False
+                    has_content_tokens = False
+                    reasoning_text = ""
+                    stream_usage = None
+                    stream_timings = None
+                    stream_done = False
+
+                    async for raw_chunk in resp.aiter_text():
+                        buffer += raw_chunk
+                        while "\n" in buffer:
+                            line, buffer = buffer.split("\n", 1)
+                            line = line.strip()
+
+                            if not line:
+                                continue
+                            if line == "data: [DONE]":
+                                if in_thinking:
+                                    if has_content_tokens:
+                                        yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
+                                    else:
+                                        yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning_text}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
+                                stream_done = True
+                                break
+                            if not line.startswith("data: "):
+                                continue
+
+                            try:
+                                data = json.loads(line[6:])
+                            except json.JSONDecodeError:
+                                continue
+
+                            _t = data.get("timings")
+                            if _t:
+                                stream_timings = _t
+                            _u = data.get("usage")
+                            if _u:
+                                stream_usage = _u
+
+                            choices = data.get("choices", [])
+                            if not choices:
+                                continue
+                            delta = choices[0].get("delta", {})
+
+                            # Handle reasoning_content -> <think> tags
+                            reasoning = delta.get("reasoning_content", "")
+                            if reasoning:
+                                reasoning_text += reasoning
+                                if not in_thinking:
+                                    in_thinking = True
+                                    yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '<think>'}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
+                                yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': reasoning}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
+
+                            # Handle content tokens
+                            token = delta.get("content", "")
+                            if token:
+                                has_content_tokens = True
+                                if in_thinking:
+                                    in_thinking = False
+                                    yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
+                                yield f"data: {json.dumps({'id': completion_id, 'object': 'chat.completion.chunk', 'created': created, 'model': model_name, 'choices': [{'index': 0, 'delta': {'content': token}, 'finish_reason': None}]}, separators = _SEP)}\n\n"
+
+                        if stream_done:
+                            break
+
+            # Final stop chunk
+            final = {
+                "id": completion_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model_name,
+                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+            }
+            yield f"data: {json.dumps(final, separators = _SEP)}\n\n"
+
+            # Usage chunk
+            if stream_usage or stream_timings:
+                usage_chunk = {
+                    "id": completion_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": model_name,
+                    "choices": [],
+                    "usage": {
+                        "prompt_tokens": (stream_usage or {}).get("prompt_tokens", 0),
+                        "completion_tokens": (stream_usage or {}).get(
+                            "completion_tokens", 0
+                        ),
+                        "total_tokens": (stream_usage or {}).get("total_tokens", 0),
+                    },
+                }
+                if stream_timings:
+                    usage_chunk["timings"] = stream_timings
+                yield f"data: {json.dumps(usage_chunk, separators = _SEP)}\n\n"
+
+            yield "data: [DONE]\n\n"
+
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
+
+    return StreamingResponse(
+        sse_generator(), media_type = "text/event-stream", headers = _SSE_HEADERS
+    )
+
+
+# ── Path B: Tool calling (asyncio.to_thread) ─────────────────
+
+
+async def _handle_tool_stream(
+    request,
+    payload,
+    llama_backend,
+    gguf_messages,
+    image_b64,
+    completion_id,
+    created,
+    model_name,
+):
+    """Handle a tool-calling streaming request via asyncio.to_thread.
+
+    Tool execution is the bottleneck, not streaming, so the thread overhead
+    is acceptable here.
+    """
+    from core.inference.tools import ALL_TOOLS
+
+    cancel_event = threading.Event()
+
+    p_enabled_tools = payload.get("enabled_tools")
+    if p_enabled_tools is not None:
+        tools_to_use = [
+            t for t in ALL_TOOLS if t["function"]["name"] in p_enabled_tools
+        ]
+    else:
+        tools_to_use = ALL_TOOLS
+
+    _sentinel = object()
+
+    async def tool_sse():
+        try:
+            first = {
+                "id": completion_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model_name,
+                "choices": [
+                    {"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}
+                ],
+            }
+            yield f"data: {json.dumps(first, separators = _SEP)}\n\n"
+
+            gen = llama_backend.generate_chat_completion_with_tools(
+                messages = gguf_messages,
+                tools = tools_to_use,
+                temperature = payload.get("temperature", 0.6),
+                top_p = payload.get("top_p", 0.95),
+                top_k = payload.get("top_k", 20),
+                min_p = payload.get("min_p", 0.01),
+                max_tokens = payload.get("max_tokens"),
+                repetition_penalty = payload.get("repetition_penalty", 1.1),
+                presence_penalty = payload.get("presence_penalty", 0.0),
+                cancel_event = cancel_event,
+                enable_thinking = payload.get("enable_thinking"),
+                auto_heal_tool_calls = payload.get("auto_heal_tool_calls", True),
+                max_tool_iterations = payload.get("max_tool_calls_per_message", 10),
+                tool_call_timeout = payload.get("tool_call_timeout", 300),
+                session_id = payload.get("session_id"),
+            )
+
+            prev_text = ""
+            _usage = None
+            _timings = None
+
+            while True:
+                if await request.is_disconnected():
+                    cancel_event.set()
+                    return
+                event = await asyncio.to_thread(next, gen, _sentinel)
+                if event is _sentinel:
+                    break
+                if event["type"] == "status":
+                    yield f"data: {json.dumps({'type': 'tool_status', 'content': event['text']})}\n\n"
+                    continue
+                if event["type"] in ("tool_start", "tool_end"):
+                    yield f"data: {json.dumps(event)}\n\n"
+                    continue
+                if event["type"] == "metadata":
+                    _usage = event.get("usage")
+                    _timings = event.get("timings")
+                    continue
+                cumulative = event.get("text", "")
+                new_text = cumulative[len(prev_text) :]
+                prev_text = cumulative
+                if not new_text:
+                    continue
+                chunk = {
+                    "id": completion_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": model_name,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": new_text},
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+                yield f"data: {json.dumps(chunk, separators = _SEP)}\n\n"
+
+            final = {
+                "id": completion_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model_name,
+                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+            }
+            yield f"data: {json.dumps(final, separators = _SEP)}\n\n"
+
+            if _usage or _timings:
+                uc = {
+                    "id": completion_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": model_name,
+                    "choices": [],
+                    "usage": {
+                        "prompt_tokens": (_usage or {}).get("prompt_tokens", 0),
+                        "completion_tokens": (_usage or {}).get("completion_tokens", 0),
+                        "total_tokens": (_usage or {}).get("total_tokens", 0),
+                    },
+                }
+                if _timings:
+                    uc["timings"] = _timings
+                yield f"data: {json.dumps(uc, separators = _SEP)}\n\n"
+
+            yield "data: [DONE]\n\n"
+        except asyncio.CancelledError:
+            cancel_event.set()
+            raise
+        except Exception as e:
+            yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
+
+    return StreamingResponse(
+        tool_sse(), media_type = "text/event-stream", headers = _SSE_HEADERS
+    )
+
+
+# ── Path C: Non-streaming ─────────────────────────────────────
+
+
+async def _handle_non_streaming(
+    payload, llama_backend, gguf_messages, image_b64, completion_id, created, model_name
+):
+    """Handle a non-streaming request. Returns a JSON response."""
+    cancel_event = threading.Event()
+
+    def _run_sync():
+        gen = llama_backend.generate_chat_completion(
+            messages = gguf_messages,
+            image_b64 = image_b64,
+            temperature = payload.get("temperature", 0.6),
+            top_p = payload.get("top_p", 0.95),
+            top_k = payload.get("top_k", 20),
+            min_p = payload.get("min_p", 0.01),
+            max_tokens = payload.get("max_tokens"),
+            repetition_penalty = payload.get("repetition_penalty", 1.0),
+            presence_penalty = payload.get("presence_penalty", 0.0),
+            stop = payload.get("stop"),
+            cancel_event = cancel_event,
+            enable_thinking = payload.get("enable_thinking"),
+        )
+        text = ""
+        usage = None
+        timings = None
+        for item in gen:
+            if isinstance(item, dict) and item.get("type") == "metadata":
+                usage = item.get("usage")
+                timings = item.get("timings")
+            elif isinstance(item, str):
+                text = item
+        return text, usage, timings
+
+    text, usage, timings = await asyncio.to_thread(_run_sync)
+
+    result = {
+        "id": completion_id,
+        "object": "chat.completion",
+        "created": created,
+        "model": model_name,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": text},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {
+            "prompt_tokens": (usage or {}).get("prompt_tokens", 0),
+            "completion_tokens": (usage or {}).get("completion_tokens", 0),
+            "total_tokens": (usage or {}).get("total_tokens", 0),
+        },
+    }
+    if timings:
+        result["timings"] = timings
+
+    return JSONResponse(
+        content = result,
+        headers = {"Access-Control-Allow-Origin": "*"},
+    )
+
+
+# ── Main endpoint ─────────────────────────────────────────────
+
+
+@stream_app.post("/stream")
+async def stream_endpoint(request: Request):
+    """
+    Stream chat completions with minimal overhead.
+
+    Three paths:
+    - Path A (hot): async httpx streaming direct to llama-server
+    - Path B: tool calling via asyncio.to_thread (tool exec is the bottleneck)
+    - Path C: non-streaming one-shot JSON response
+    """
+    (
+        payload,
+        llama_backend,
+        gguf_messages,
+        image_b64,
+        completion_id,
+        created,
+        model_name,
+    ) = await _validate_request(request)
+
+    # Path C: Non-streaming
+    stream = payload.get("stream", True)
+    if not stream:
+        return await _handle_non_streaming(
+            payload,
+            llama_backend,
+            gguf_messages,
+            image_b64,
+            completion_id,
+            created,
+            model_name,
+        )
+
+    # Path B: Tool calling
+    use_tools = payload.get("use_tools", False)
+    if use_tools and llama_backend.supports_tools:
+        return await _handle_tool_stream(
+            request,
+            payload,
+            llama_backend,
+            gguf_messages,
+            image_b64,
+            completion_id,
+            created,
+            model_name,
+        )
+
+    # Path A: Direct async streaming (hot path)
+    return await _handle_async_stream(
+        request,
+        payload,
+        llama_backend,
+        gguf_messages,
+        image_b64,
+        completion_id,
+        created,
+        model_name,
+    )
+
+
+# ── Server lifecycle ──────────────────────────────────────────
+
+
+def start_streaming_server(port: int) -> None:
+    """Start the streaming server in the current thread (blocking). Use in a daemon thread."""
+    import uvicorn
+
+    uvicorn.run(
+        stream_app,
+        host = "127.0.0.1",
+        port = port,
+        log_level = "warning",
+        access_log = False,
+    )
+
+
+def find_free_port() -> int:
+    """Find a free TCP port."""
+    import socket
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]

studio/frontend/src/features/chat/api/chat-api.ts

@@ -177,12 +177,46 @@ export async function* streamChatCompletions(
   payload: OpenAIChatCompletionsRequest,
   signal: AbortSignal,
 ): AsyncGenerator<OpenAIChatChunk> {
-  const response = await authFetch("/v1/chat/completions", {
+  // Try to acquire a fast-path token for the dedicated streaming server.
-    method: "POST",
+  // Falls back silently to the baseline endpoint on any failure.
-    headers: { "Content-Type": "application/json" },
+  let useFastPath = false;
-    body: JSON.stringify(payload),
+  let fastUrl = "";
-    signal,
+  let fastToken = "";
-  });
+
+  try {
+    const streamUrlResp = await authFetch("/api/inference/stream-url", { signal });
+    if (streamUrlResp.ok) {
+      const info = (await streamUrlResp.json()) as {
+        supported?: boolean;
+        stream_url?: string;
+        token?: string;
+      };
+      if (info.supported && info.stream_url && info.token) {
+        fastUrl = info.stream_url;
+        fastToken = info.token;
+        useFastPath = true;
+      }
+    }
+  } catch {
+    /* fall back silently to baseline */
+  }
+
+  const response = useFastPath
+    ? await fetch(fastUrl, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          "X-Stream-Token": fastToken,
+        },
+        body: JSON.stringify(payload),
+        signal,
+      })
+    : await authFetch("/v1/chat/completions", {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(payload),
+        signal,
+      });
 
   if (!response.ok) {
     const body = await response.json().catch(() => null);