@@ -59,10 +97,10 @@ def show_link(port: int = 8888):
height="48" style="display:block;">
Unsloth Studio is Ready!
-
+ font-weight: 800; font-size: 16px; cursor: pointer;">
Open Unsloth Studio
@@ -77,6 +115,75 @@ def show_link(port: int = 8888):
display(HTML(html))
+def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
+ """Return True if a Studio backend is already answering health checks on *port*."""
+ import urllib.request
+
+ try:
+ with urllib.request.urlopen(
+ f"http://localhost:{port}/api/health", timeout = timeout
+ ):
+ return True
+ except Exception:
+ return False
+
+
+def _show_and_embed(port: int):
+ """Embed the Studio inline for *port* with a branded header bar.
+
+ Fetches the Colab proxy URL once (registering the port with Colab's
+ reverse-proxy at the same time) then renders a header bar + full-height
+ iframe as a single HTML block.
+
+ Falls back to ``serve_kernel_port_as_iframe`` if ``IPython.display.HTML``
+ is unavailable for any reason.
+ """
+ url = get_colab_url(port)
+ logger.info(f"🌐 Unsloth Studio URL: {url}")
+
+ try:
+ from IPython.display import HTML, display
+
+ iframe_id = f"unsloth-studio-{port}"
+
+ # Truncated URL shown in the header — best-effort, falls back to full URL.
+ try:
+ port_prefix = f"{port}-"
+ idx = url.index(port_prefix)
+ next_dash = url.index("-", idx + len(port_prefix))
+ short_url = url[: next_dash + 1] + "..."
+ except (ValueError, IndexError):
+ short_url = url
+
+ display(
+ HTML(f"""
+
+""")
+ )
+ except Exception:
+ # Fallback: Colab's built-in helper (less control, but always works)
+ try:
+ from google.colab import output as colab_output
+
+ colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%")
+ except ImportError:
+ pass
+
+
def start(port: int = 8888):
"""
Start Unsloth Studio server in Colab and display the URL.
@@ -85,10 +192,26 @@ def start(port: int = 8888):
from colab import start
start()
"""
- import sys
+ import time
logger.info("🦥 Starting Unsloth Studio...")
+ # --- Fast path: Studio is already running (cell re-run) ---
+ # Re-launching would either collide on the port or silently shift to a new
+ # port and confuse the user. Just re-show the link and iframe instead.
+ if _is_studio_healthy(port):
+ logger.info(
+ f" Studio is already running on port {port} — reusing existing server."
+ )
+ _show_and_embed(port)
+ try:
+ for _ in range(10000):
+ time.sleep(300)
+ print("=", end = "", flush = True)
+ except KeyboardInterrupt:
+ logger.info("\nUnsloth Studio keepalive stopped.")
+ return
+
logger.info(" Loading backend...")
from run import run_server
@@ -96,18 +219,63 @@ def start(port: int = 8888):
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
- if not frontend_path.exists():
+ if not (frontend_path / "index.html").exists():
logger.info("❌ Frontend not built! Please run the setup cell first.")
return
logger.info(" Starting server...")
- # Start server silently
- run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
+ try:
+ app = run_server(
+ host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True
+ )
+ except SystemExit as exc:
+ logger.error(f"❌ Unsloth Studio failed to start: {exc}")
+ return
+ except Exception as exc:
+ logger.error(f"❌ Unsloth Studio failed to start: {exc}")
+ return
- logger.info(" Server started!")
+ # run_server auto-increments the port when the requested one is already in
+ # use (e.g. Jupyter occupying 8888). Read back the actual bound port so the
+ # Colab proxy URL and iframe always point at the right place.
+ actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port
- # Show the clickable link with real URL
- show_link(port)
+ logger.info(f" Server started on port {actual_port}!")
+
+ # Poll health endpoint to confirm the server is truly reachable before
+ # showing the link and registering the iframe — avoids the race where
+ # ready_event fires but the process hasn't finished binding.
+ import urllib.request
+
+ server_ready = False
+ for _ in range(40):
+ try:
+ with urllib.request.urlopen(
+ f"http://localhost:{actual_port}/api/health", timeout = 1
+ ):
+ server_ready = True
+ break
+ except Exception:
+ time.sleep(0.5)
+
+ if not server_ready:
+ logger.error(
+ f"❌ Unsloth Studio did not become healthy on port {actual_port}. "
+ "Check for errors above."
+ )
+ return
+
+ _show_and_embed(actual_port)
+
+ # Keep kernel alive so the daemon server thread stays running.
+ # Handle KeyboardInterrupt cleanly so the user gets a readable message
+ # rather than a raw traceback when they interrupt the cell.
+ try:
+ for _ in range(10000):
+ time.sleep(300)
+ print("=", end = "", flush = True)
+ except KeyboardInterrupt:
+ logger.info("\nUnsloth Studio keepalive stopped.")
if __name__ == "__main__":
diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py
new file mode 100644
index 0000000000..5650a60ee2
--- /dev/null
+++ b/studio/backend/core/_torchao_stub.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Shared torchao Windows-ROCm import stub.
+
+torchao (pulled in by transformers.quantizers) imports
+torch.distributed._functional_collectives at module level, which imports
+distributed_c10d.py unconditionally — that file crashes on Windows ROCm because
+torch._C._distributed_c10d (the RCCL backend) is absent.
+torch/distributed/__init__.py itself is guarded by `if is_available()` so
+`import torch.distributed` alone is safe; the crash only comes via torchao's
+import chain. Stubbing torchao short-circuits it entirely.
+_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports.
+
+This logic used to be duplicated inline inside run_export_process() and
+run_training_process(); it now lives here so both worker subprocesses call the
+single `install_torchao_windows_rocm_stub()` entrypoint before importing
+transformers / unsloth_zoo.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+import importlib.abc
+import importlib.machinery
+
+_STUB_SENTINEL = object()
+
+
+# Metaclass for stub types so that isinstance(x, StubClass) returns False
+# instead of raising TypeError ("arg 2 must be a type").
+# peft/tuners/lora/torchao.py does:
+# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor
+# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor))
+# If those names resolve to stub modules rather than types, isinstance() raises.
+class _StubTypeMeta(type):
+ def __instancecheck__(cls, instance):
+ return False
+
+ def __subclasscheck__(cls, subclass):
+ return False
+
+ def __getattr__(cls, attr):
+ if attr.startswith("__"):
+ raise AttributeError(attr)
+ child = _StubTypeMeta(attr, (), {})
+ setattr(cls, attr, child)
+ return child
+
+ def __call__(cls, *args, **kwargs):
+ return None
+
+
+def _make_stub_type(name):
+ """Stub class: accepted by isinstance() (always False), supports attr access."""
+ return _StubTypeMeta(name, (), {})
+
+
+def _make_mod_stub(mod_name):
+ m = types.ModuleType(mod_name)
+ m.__path__ = []
+ m.__package__ = mod_name
+ m._unsloth_stub = _STUB_SENTINEL
+ m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True)
+
+ def _ga(attr, _m = m, _n = mod_name):
+ if attr.startswith("__"):
+ raise AttributeError(attr)
+ # Return a stub CLASS (not a module) so that isinstance(x, attr)
+ # works and returns False instead of raising TypeError.
+ child = _make_stub_type(f"{_n}.{attr}")
+ setattr(_m, attr, child)
+ return child
+
+ m.__getattr__ = _ga
+ return m
+
+
+class _StubSubpackageLoader(importlib.abc.Loader):
+ def __init__(self, mod_name):
+ self._mod_name = mod_name
+
+ def create_module(self, spec):
+ return _make_mod_stub(self._mod_name)
+
+ def exec_module(self, module):
+ pass
+
+
+class _StubSubpackageFinder(importlib.abc.MetaPathFinder):
+ def find_spec(self, fullname, path, target = None):
+ if "." not in fullname:
+ return None
+ parent = sys.modules.get(fullname.rsplit(".", 1)[0])
+ if parent is None:
+ return None
+ if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
+ return None
+ return importlib.machinery.ModuleSpec(
+ fullname, _StubSubpackageLoader(fullname), is_package = True
+ )
+
+
+def install_torchao_windows_rocm_stub() -> None:
+ """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash.
+
+ No-op on every other platform (Windows CUDA included — there torchao is real
+ and shadowing it would break torchao-based quantization paths). Must run
+ before any import of transformers / unsloth_zoo. Safe to call once per worker
+ process.
+ """
+ # Gate on the active torch runtime, not env-var presence -- HIP_PATH /
+ # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a
+ # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip
+ # but still encode "rocm" in torch.__version__, so accept either.
+ _is_win32_rocm = False
+ if sys.platform == "win32":
+ try:
+ import torch as _torch_probe
+
+ _is_win32_rocm = bool(
+ getattr(getattr(_torch_probe, "version", None), "hip", None)
+ or "rocm" in getattr(_torch_probe, "__version__", "").lower()
+ )
+ del _torch_probe
+ except Exception:
+ pass
+ if _is_win32_rocm:
+ # Register the finder only on Windows ROCm -- on other platforms there
+ # are no stub modules seeded, so appending is a pure accumulation.
+ sys.meta_path.append(_StubSubpackageFinder())
+ # Seed torchao top-level + key submodules; the finder handles the rest.
+ for _tao_name in (
+ "torchao",
+ "torchao.quantization",
+ "torchao.dtypes",
+ "torchao.float8",
+ "torchao.utils",
+ ):
+ if _tao_name not in sys.modules:
+ sys.modules[_tao_name] = _make_mod_stub(_tao_name)
diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py
index b4ec0ccd94..85b567885f 100644
--- a/studio/backend/core/data_recipe/service.py
+++ b/studio/backend/core/data_recipe/service.py
@@ -176,12 +176,21 @@ def build_mcp_providers(
) -> list:
from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports]
+ # Same gate as the chat MCP path: stdio providers spawn a local subprocess,
+ # so only build them when this host allows it (desktop / explicit opt-in).
+ # Skip them otherwise so a recipe carried onto a hosted host cannot spawn.
+ from core.inference.mcp_client import stdio_mcp_enabled
+
+ stdio_allowed = stdio_mcp_enabled()
+
providers: list[MCPProvider | LocalStdioMCPProvider] = []
for provider in recipe.get("mcp_providers", []):
if not isinstance(provider, dict):
continue
provider_type = provider.get("provider_type")
if provider_type == "stdio":
+ if not stdio_allowed:
+ continue
env = provider.get("env")
if not isinstance(env, dict):
env = {}
diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py
index f77b1966c4..defcff924b 100644
--- a/studio/backend/core/export/worker.py
+++ b/studio/backend/core/export/worker.py
@@ -439,6 +439,15 @@ def run_export_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
+ # ── 1c. Stub torchao on Windows ROCm ──
+ # Shared with the training worker; see core/_torchao_stub.py for the full
+ # rationale (torchao -> torch.distributed._functional_collectives crashes on
+ # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm.
+ # Must run before any import of transformers / unsloth_zoo.
+ from core._torchao_stub import install_torchao_windows_rocm_stub
+
+ install_torchao_windows_rocm_stub()
+
# ── 2. Import ML libraries (fresh in this clean process) ──
try:
_send_response(
diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py
index 8f34bb23fc..8a1edd608b 100644
--- a/studio/backend/core/inference/external_provider.py
+++ b/studio/backend/core/inference/external_provider.py
@@ -8,10 +8,12 @@ Most registry providers expose OpenAI-compatible /v1/chat/completions endpoints;
Anthropic uses native Messages API with translation in this client.
"""
+import base64
import json as _json
+import mimetypes
import re
import time
-from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional
+from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional, Union
from urllib.parse import urlparse
import httpx
@@ -506,6 +508,250 @@ def _apply_mistral_reasoning_controls(
_http_client = httpx.AsyncClient()
+# Cap per-image fetch well below Gemini's ~20 MB total request budget.
+_GEMINI_REMOTE_IMAGE_MAX_BYTES = 10 * 1024 * 1024
+_GEMINI_REMOTE_IMAGE_TIMEOUT_S = 15.0
+
+
+def _safe_fetch_image_for_gemini_sync(
+ url: str,
+ fallback_mime: str,
+ max_bytes: int = _GEMINI_REMOTE_IMAGE_MAX_BYTES,
+) -> Optional[tuple[str, str]]:
+ """Synchronous IP-pinned HTTPS image fetch with SSRF guards.
+
+ Uses the same pinned-IP + SNI pattern as `tools._fetch_page_text` so
+ DNS rebinding between validation and the actual connection cannot
+ redirect us to a private/metadata address. Follows up to 4 hops,
+ re-validating each redirect target. Returns (mime, base64) or None.
+
+ `max_bytes` is clamped to the per-image cap and additionally lets
+ the caller pass the remaining per-request budget so an over-budget
+ URL is rejected via Content-Length (or read short-circuit) instead
+ of being fully downloaded then discarded after the fact.
+ """
+ import urllib.error
+ import urllib.request
+ from urllib.parse import urljoin, urlunparse
+
+ # Refuse upfront if the per-request budget is already spent.
+ _byte_limit = min(max(0, int(max_bytes)), _GEMINI_REMOTE_IMAGE_MAX_BYTES)
+ if _byte_limit <= 0:
+ return None
+
+ # Share tools.py's pinned-IP hardening: validate-once-then-pin.
+ from .tools import (
+ _NoRedirect,
+ _SNIHTTPSHandler,
+ _validate_and_resolve_host,
+ )
+
+ def _safe_parse_https(raw_url: str) -> Optional[tuple[Any, str, int]]:
+ """Validate https + hostname + port. Returns (parsed, host, port) or
+ None. Handles malformed-port and malformed-bracketed-IPv6 URLs that
+ would otherwise raise ValueError mid-build.
+ """
+ try:
+ parsed_url = urlparse(raw_url)
+ host_value = parsed_url.hostname
+ port_value = parsed_url.port or 443
+ except (ValueError, UnicodeError) as _err:
+ logger.info(
+ "Gemini image fetch: refusing malformed url err=%s",
+ type(_err).__name__,
+ )
+ return None
+ scheme_value = (parsed_url.scheme or "").lower()
+ if scheme_value != "https":
+ logger.info(
+ "Gemini image fetch: refusing non-https scheme=%s",
+ scheme_value,
+ )
+ return None
+ if not host_value:
+ logger.info("Gemini image fetch: refusing url with no hostname")
+ return None
+ return parsed_url, host_value, port_value
+
+ parsed_info = _safe_parse_https(url)
+ if parsed_info is None:
+ return None
+ parsed, current_host, current_port = parsed_info
+ current_url = url
+ ok, reason, pinned_ip = _validate_and_resolve_host(current_host, current_port)
+ if not ok:
+ logger.warning(
+ "Gemini image fetch: refusing host=%s reason=%s",
+ current_host,
+ reason,
+ )
+ return None
+
+ for _hop in range(4):
+ # Pin to validated IP; SNI + cert still use hostname via _SNIHTTPSHandler.
+ cp_info = _safe_parse_https(current_url)
+ if cp_info is None:
+ return None
+ cp, _cp_host, _cp_port = cp_info
+ ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip
+ ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str
+ pinned_url = urlunparse(cp._replace(netloc = ip_netloc))
+
+ opener = urllib.request.build_opener(
+ _NoRedirect,
+ _SNIHTTPSHandler(current_host),
+ )
+ req = urllib.request.Request(
+ pinned_url,
+ headers = {"Host": current_host},
+ method = "GET",
+ )
+
+ try:
+ resp = opener.open(req, timeout = _GEMINI_REMOTE_IMAGE_TIMEOUT_S)
+ except urllib.error.HTTPError as e:
+ if e.code not in (301, 302, 303, 307, 308):
+ logger.info(
+ "Gemini image fetch: status=%d host=%s",
+ e.code,
+ current_host,
+ )
+ return None
+ location = e.headers.get("Location")
+ if not location:
+ return None
+ try:
+ current_url = urljoin(current_url, location)
+ except (ValueError, UnicodeError) as _err:
+ logger.info(
+ "Gemini image fetch: refusing malformed redirect err=%s",
+ type(_err).__name__,
+ )
+ return None
+ rp_info = _safe_parse_https(current_url)
+ if rp_info is None:
+ return None
+ _rp, current_host, current_port = rp_info
+ ok2, reason2, pinned_ip = _validate_and_resolve_host(
+ current_host, current_port
+ )
+ if not ok2:
+ logger.warning(
+ "Gemini image fetch: refusing redirect host=%s reason=%s",
+ current_host,
+ reason2,
+ )
+ return None
+ continue
+ except (urllib.error.URLError, OSError) as _err:
+ logger.warning(
+ "Gemini image fetch failed host=%s err=%s",
+ current_host,
+ type(_err).__name__,
+ )
+ return None
+
+ with resp:
+ status = getattr(resp, "status", None) or resp.getcode()
+ if status != 200:
+ logger.info(
+ "Gemini image fetch: status=%s host=%s", status, current_host
+ )
+ return None
+ _hdr_mime = (
+ (resp.headers.get("content-type") or "").split(";")[0].strip().lower()
+ )
+ # Declared non-image MIME is a refusal; missing MIME falls back to caller's.
+ if _hdr_mime and not _hdr_mime.startswith("image/"):
+ logger.info(
+ "Gemini image fetch: non-image content-type=%s host=%s",
+ _hdr_mime,
+ current_host,
+ )
+ return None
+ _final_mime_pre = _hdr_mime if _hdr_mime else fallback_mime
+ if not isinstance(_final_mime_pre, str) or not _final_mime_pre.startswith(
+ "image/"
+ ):
+ logger.info(
+ "Gemini image fetch: missing content-type and no image fallback host=%s",
+ current_host,
+ )
+ return None
+ _hdr_len = resp.headers.get("content-length")
+ if _hdr_len and _hdr_len.isdigit() and int(_hdr_len) > _byte_limit:
+ logger.info(
+ "Gemini image fetch: declared %s bytes exceeds cap=%s host=%s",
+ _hdr_len,
+ _byte_limit,
+ current_host,
+ )
+ return None
+ # Read cap+1 to detect oversize without buffering unbounded data.
+ raw = resp.read(_byte_limit + 1)
+ if len(raw) > _byte_limit:
+ logger.info(
+ "Gemini image fetch: streamed bytes exceed cap=%s host=%s",
+ _byte_limit,
+ current_host,
+ )
+ return None
+ return _final_mime_pre, base64.b64encode(raw).decode("ascii")
+
+ logger.info("Gemini image fetch: too many redirects host=%s", current_host)
+ return None
+
+
+async def _safe_fetch_image_for_gemini(
+ url: str,
+ fallback_mime: str,
+ max_bytes: int = _GEMINI_REMOTE_IMAGE_MAX_BYTES,
+) -> Optional[tuple[str, str]]:
+ """Async wrapper running the IP-pinned fetch on a worker thread.
+
+ SSRF guards (https only, pinned IP, per-hop redirect re-check, size
+ cap, image/* content-type) live in the sync helper. `max_bytes`
+ carries the remaining per-request budget so over-budget URLs are
+ rejected up front.
+ """
+ import asyncio
+
+ return await asyncio.to_thread(
+ _safe_fetch_image_for_gemini_sync, url, fallback_mime, max_bytes
+ )
+
+
+# Synthetic-tool names stamped onto outbound _toolEvent.arguments so the
+# frontend can distinguish provider-side cards from real user-declared
+# tools of the same name. Mirrored on the TS side.
+_SERVER_SIDE_BUILTIN_TOOL_NAMES = frozenset(
+ {"web_search", "web_fetch", "code_execution", "image_generation"}
+)
+
+
+def _stamp_server_tool_marker(payload: dict[str, Any]) -> None:
+ """Tag synthetic provider-side tool events so the frontend can
+ distinguish them from real user-declared / local function tools of
+ the same name. The marker rides on `arguments._server_tool` and is
+ only added for known server-side builtin names; user-supplied
+ tool calls echoed back through these helpers (e.g. Kimi
+ `$web_search`) keep their existing shape because we keep this scoped
+ to the canonical builtin names.
+ """
+ if not isinstance(payload, dict):
+ return
+ if payload.get("type") != "tool_start":
+ return
+ name = payload.get("tool_name")
+ if not isinstance(name, str) or name not in _SERVER_SIDE_BUILTIN_TOOL_NAMES:
+ return
+ args = payload.get("arguments")
+ if not isinstance(args, dict):
+ args = {}
+ payload["arguments"] = args
+ args["_server_tool"] = True
+
+
def _build_kimi_tool_end(
synthetic_chunk_fn: Any,
tool_call_id: str,
@@ -546,16 +792,23 @@ class ExternalProviderClient:
):
self.provider_type = provider_type
self.base_url = base_url.rstrip("/")
+ # Strip a legacy `/openai` suffix from Google-hosted bases so
+ # configs saved before the native switch still route correctly.
+ # Custom proxy paths ending in `/openai` are left untouched.
+ if self.provider_type == "gemini":
+ _parsed_base = urlparse(self.base_url)
+ if (
+ (_parsed_base.hostname or "").lower()
+ == "generativelanguage.googleapis.com"
+ and _parsed_base.path.rstrip("/") == "/v1beta/openai"
+ ):
+ self.base_url = self.base_url[: -len("/openai")]
self.api_key = api_key
self._timeout = httpx.Timeout(timeout, connect = 10.0)
- # Separate timeout for SSE streams: reasoning-heavy providers
- # (Anthropic Opus 4.7 with adaptive thinking, OpenAI gpt-5.x via
- # /v1/responses) can pause for tens of seconds between bytes
- # while the model is internally thinking. httpx's read timeout is
- # the *gap* between successive reads, not a wall clock — so
- # disabling it lets long thinks complete without cutting the
- # stream prematurely. connect/write/pool keep the 10s / 120s
- # bounds so genuine network failures still surface.
+ # Disable read timeout on SSE streams: reasoning-heavy models
+ # pause tens of seconds between bytes while thinking, and httpx's
+ # read timeout is the per-byte gap, not wall clock. connect/write
+ # bounds still surface real network failures.
self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None)
def _auth_headers(self) -> dict[str, str]:
@@ -566,6 +819,14 @@ class ExternalProviderClient:
auth_header = provider_info.get("auth_header", "Authorization")
auth_prefix = provider_info.get("auth_prefix", "Bearer ")
+ # Non-Google Gemini bases (LiteLLM, custom gateways) use OAI-compat
+ # Bearer auth, not Google's x-goog-api-key. Override the registry default.
+ if self.provider_type == "gemini":
+ _host = (urlparse(self.base_url).hostname or "").lower()
+ if _host != "generativelanguage.googleapis.com":
+ auth_header = "Authorization"
+ auth_prefix = "Bearer "
+
headers = {"Content-Type": "application/json"}
# Skip auth header when api_key is empty (optional for local providers);
# httpx rejects an empty `Bearer ` value as "Illegal header value".
@@ -580,6 +841,12 @@ class ExternalProviderClient:
from core.inference.providers import get_provider_info
info = get_provider_info(self.provider_type) or {}
+ # Google-hosted Gemini uses the native translator; non-Google
+ # bases stay on OAI-compat so LiteLLM / custom proxies still work.
+ if self.provider_type == "gemini":
+ _host = (urlparse(self.base_url).hostname or "").lower()
+ if _host != "generativelanguage.googleapis.com":
+ return True
return info.get("openai_compatible", True)
async def stream_chat_completion(
@@ -594,11 +861,13 @@ class ExternalProviderClient:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
enabled_tools: Optional[list[str]] = None,
- enable_prompt_caching: Optional[bool] = None,
+ enable_prompt_caching: Optional[Union[bool, str]] = None,
openai_code_exec_container_id: Optional[str] = None,
anthropic_code_exec_container_id: Optional[str] = None,
prompt_cache_ttl: Optional[str] = None,
compaction_threshold: Optional[int] = None,
+ tools: Optional[list[dict[str, Any]]] = None,
+ tool_choice: Optional[Any] = None,
fast_mode: Optional[bool] = None,
stream: bool = True,
) -> AsyncGenerator[str, None]:
@@ -616,7 +885,35 @@ class ExternalProviderClient:
``fast_mode`` only applies to Anthropic Opus 4.6 / 4.7 (silently
dropped elsewhere); adds the beta header and ``speed: "fast"``.
"""
+ # tool_choice="none" hard-disables hosted/builtin tools across
+ # every provider so enabled_tools cannot accidentally bill or leak.
+ tool_choice_disabled = (
+ isinstance(tool_choice, str) and tool_choice.strip().lower() == "none"
+ )
+
if not self._is_openai_compatible():
+ # Gemini speaks its own native REST shape (contents/parts);
+ # `_stream_gemini` translates request/response into the OpenAI
+ # Chat Completions chunk format the rest of Studio expects.
+ # API reference: https://ai.google.dev/gemini-api/docs
+ if self.provider_type == "gemini":
+ async for line in self._stream_gemini(
+ messages,
+ model,
+ temperature,
+ top_p,
+ max_tokens,
+ top_k,
+ presence_penalty,
+ enabled_tools,
+ enable_prompt_caching,
+ enable_thinking,
+ reasoning_effort,
+ tools,
+ tool_choice,
+ ):
+ yield line
+ return
async for line in self._stream_anthropic(
messages,
model,
@@ -631,6 +928,7 @@ class ExternalProviderClient:
anthropic_code_exec_container_id,
prompt_cache_ttl,
compaction_threshold,
+ tool_choice,
fast_mode = fast_mode,
):
yield line
@@ -654,20 +952,25 @@ class ExternalProviderClient:
enable_prompt_caching,
openai_code_exec_container_id,
compaction_threshold,
+ tools,
+ tool_choice,
):
yield line
return
- # Kimi's $web_search is a builtin_function that requires a client
- # round-trip: the first call returns a tool_calls envelope with
- # function.arguments populated; the caller echoes those arguments
- # back as a role=tool message; the second call streams the final
- # answer with the search incorporated. The doc also mandates
- # disabling thinking while $web_search is active. Route to a
- # dedicated helper so the default OAI-compat path stays single-pass.
- # https://platform.kimi.ai/docs/guide/use-web-search
+ # Kimi $web_search needs a 2-call round-trip + thinking off; route
+ # to a helper. Forced-function tool_choice suppresses it.
+ # https://platform.kimi.ai/docs/guide/use-web-search
+ _kimi_tool_choice_forced_function = (
+ isinstance(tool_choice, dict)
+ and tool_choice.get("type") == "function"
+ and isinstance(tool_choice.get("function"), dict)
+ and bool(tool_choice["function"].get("name"))
+ )
if (
self.provider_type == "kimi"
+ and not tool_choice_disabled
+ and not _kimi_tool_choice_forced_function
and enabled_tools
and "web_search" in enabled_tools
):
@@ -694,28 +997,18 @@ class ExternalProviderClient:
else:
body["max_tokens"] = max_tokens
- # Strip body fields a provider's registry entry declares unusable —
- # reasoning-class models that lock these to fixed defaults (e.g.
- # Kimi k2.5/k2.6 only accept temperature=1, top_p=1) 400 otherwise.
- # The frontend capability map already hides the matching sliders;
- # this is the matching guard for the pydantic default that the
- # route layer would otherwise still fill in.
+ # Drop fields the registry flags as unusable so reasoning-class
+ # models with fixed defaults (Kimi k2.6 etc) don't 400 on pydantic
+ # default values that the route layer still fills in.
from core.inference.providers import get_provider_info
provider_info = get_provider_info(self.provider_type) or {}
for field in provider_info.get("body_omit", ()):
body.pop(field, None)
- # Kimi (kimi-k2.6, kimi-k2-thinking) accepts a boolean thinking toggle
- # via a top-level `thinking` field (the docs show it nested under
- # extra_body, but that is an OpenAI Python SDK convention; on the
- # wire it merges into the request body).
- # - kimi-k2.6 defaults to thinking enabled; clients can pass
- # {"type": "disabled"} to suppress it.
- # - kimi-k2-thinking is always on; we never send disabled there.
- # `keep: all` retains every thinking chunk through the stream, which
- # is what we need so our frontend can wrap reasoning_content into
- # the chat reasoning panel.
+ # Kimi thinking is a top-level body field. kimi-k2-thinking is
+ # always on (ignore the toggle); kimi-k2.6 defaults on, can be
+ # disabled. `keep: all` preserves every chunk for the UI panel.
if self.provider_type == "kimi" and enable_thinking is not None:
if model == "kimi-k2-thinking":
# Always on; ignore client toggle to avoid an API-level reject.
@@ -736,17 +1029,9 @@ class ExternalProviderClient:
tpl_kw["enable_thinking"] = bool(enable_thinking)
body["chat_template_kwargs"] = tpl_kw
- # OpenRouter exposes a unified `reasoning` parameter on every
- # chat-completion request — the gateway routes it to whichever
- # underlying model actually supports reasoning, and silently
- # no-ops for ones that don't. Documented at
- # https://openrouter.ai/docs/guides/best-practices/reasoning-tokens
- # Shape: `reasoning: {enabled?: bool, effort?: low|medium|high,
- # max_tokens?: N, exclude?: bool}` with effort and max_tokens
- # mutually exclusive. We forward either an effort level (when
- # the user picked one) or a bare {enabled: true}. A small set of
- # known routes rejects explicit disable with 400 ("Reasoning is
- # mandatory for this endpoint ..."), so only those omit "off".
+ # OpenRouter's unified `reasoning` field gates per-model thinking.
+ # Some routes (`*_MANDATORY_REASONING_MODELS`) 400 on explicit off.
+ # https://openrouter.ai/docs/guides/best-practices/reasoning-tokens
if self.provider_type == "openrouter":
normalized_or_model = model.strip().lower()
if reasoning_effort in ("low", "medium", "high"):
@@ -759,17 +1044,22 @@ class ExternalProviderClient:
else:
body["reasoning"] = {"enabled": False}
- # OpenRouter web-search plugin — universal shape that works
- # for every model id, including the `openrouter/free` and
- # `openrouter/auto` meta-routers. Documented at
- # https://openrouter.ai/docs/guides/features/plugins/web-search
- # The `:online` model-suffix shortcut is "exactly equivalent
- # to" this plugin per the same doc, but only works on
- # concrete model ids — meta-routers reject the suffix.
- # `plugins: [{id: "web"}]` works everywhere, no model id
- # rewrite needed, and idempotent if some future call site
- # adds the entry first.
- if enabled_tools and "web_search" in enabled_tools:
+ # OpenRouter web plugin works on every model id including
+ # meta-routers (unlike the `:online` suffix). Forced-function
+ # tool_choice suppresses it, matching Gemini/Anthropic.
+ # https://openrouter.ai/docs/guides/features/plugins/web-search
+ _or_tool_choice_forced_function = (
+ isinstance(tool_choice, dict)
+ and tool_choice.get("type") == "function"
+ and isinstance(tool_choice.get("function"), dict)
+ and bool(tool_choice["function"].get("name"))
+ )
+ if (
+ not tool_choice_disabled
+ and not _or_tool_choice_forced_function
+ and enabled_tools
+ and "web_search" in enabled_tools
+ ):
plugins = list(body.get("plugins") or [])
if not any(
isinstance(p, dict) and p.get("id") == "web" for p in plugins
@@ -781,6 +1071,15 @@ class ExternalProviderClient:
body.get("model"),
)
+ # Forward OpenAI-style function tools / tool_choice on every
+ # OAI-compat route (incl. custom Gemini OpenAI proxies like
+ # LiteLLM). Without this, callers that wire user-defined tools
+ # silently lose function-calling on non-native providers.
+ if tools:
+ body["tools"] = tools
+ if tool_choice is not None:
+ body["tool_choice"] = tool_choice
+
url = f"{self.base_url}/chat/completions"
logger.info(
"Proxying chat completion to %s (provider=%s, model=%s)",
@@ -816,30 +1115,22 @@ class ExternalProviderClient:
)
return
- # NOTE: manual __anext__ loop instead of `async for` is intentional.
- # On Python 3.13 + httpcore 1.0.x, `async for` auto-calls aclose() on
- # early exit (break/return/GeneratorExit) BEFORE our finally block runs.
- # That propagates GeneratorExit into PoolByteStream.__aiter__() while it
- # calls `await self.aclose()` inside `with AsyncShieldCancellation()`,
- # triggering "RuntimeError: async generator ignored GeneratorExit".
- # Fix: call response.aclose() FIRST (sets PoolByteStream._closed=True),
- # then lines_gen.aclose() is a no-op and GeneratorExit re-raises cleanly.
+ # Manual __anext__ (not `async for`) so we can close
+ # the response BEFORE lines_gen, avoiding the httpcore
+ # 1.0 GeneratorExit -> RuntimeError path on Python 3.13.
lines_gen = response.aiter_lines().__aiter__()
- # Best-effort diagnostics for the default OAI-compat path. Without
- # this, OpenRouter mid-stream errors (200 OK + error event in the
- # SSE body) and OpenRouter-router model selection were invisible
- # in the backend logs — the user only saw "Provider returned
- # error" in the UI with no trail on the server side.
+ # Diagnostic counters for the OAI-compat path; surfaces
+ # OpenRouter mid-stream errors that would otherwise be
+ # invisible server-side.
event_counts: dict[str, int] = {}
chosen_model: Optional[str] = None
- # Web-search tool-card synthesis for OpenRouter. The gateway
- # doesn't emit structured web_search_call events — citations
- # come back as `annotations` of type=url_citation on delta /
- # message objects. Mirror the OpenAI/Anthropic UX by yielding
- # a synthetic tool_start at stream open and tool_end at
- # stream close with the collected citation list.
+ # OpenRouter has no web_search_call events — citations
+ # arrive as url_citation annotations. Synthesise a
+ # tool_start/tool_end pair to match the OpenAI/Anthropic UX.
web_search_active = (
self.provider_type == "openrouter"
+ and not tool_choice_disabled
+ and not _or_tool_choice_forced_function
and bool(enabled_tools)
and "web_search" in (enabled_tools or [])
)
@@ -849,6 +1140,7 @@ class ExternalProviderClient:
web_search_tool_ended = False
def _emit_synthetic_tool_event(payload: dict[str, Any]) -> str:
+ _stamp_server_tool_marker(payload)
chunk = {
"id": f"chatcmpl-{self.provider_type}-synthetic",
"object": "chat.completion.chunk",
@@ -1104,6 +1396,7 @@ class ExternalProviderClient:
synthetic_id = f"chatcmpl-{self.provider_type}-synthetic"
def _synthetic_chunk(payload: dict[str, Any]) -> str:
+ _stamp_server_tool_marker(payload)
chunk = {
"id": synthetic_id,
"object": "chat.completion.chunk",
@@ -1441,6 +1734,7 @@ class ExternalProviderClient:
anthropic_code_exec_container_id: Optional[str] = None,
prompt_cache_ttl: Optional[str] = None,
compaction_threshold: Optional[int] = None,
+ tool_choice: Optional[Any] = None,
*,
fast_mode: Optional[bool] = None,
) -> AsyncGenerator[str, None]:
@@ -1471,6 +1765,42 @@ class ExternalProviderClient:
continue
content = msg.get("content")
+ # OpenAI role="tool" with list content -> Anthropic native
+ # tool_result block on a user message. Translating in the
+ # string-content branch only (below) leaves the list-content
+ # form forwarded as an invalid `role:"tool"` message that
+ # Anthropic rejects. Handle both upfront.
+ if msg.get("role") == "tool":
+ _tr_id = msg.get("tool_call_id") or ""
+ if isinstance(content, list):
+ _flat_parts: list[str] = []
+ for part in content:
+ if (
+ isinstance(part, dict)
+ and part.get("type") == "text"
+ and part.get("text")
+ ):
+ _flat_parts.append(str(part["text"]))
+ _flat_result = "".join(_flat_parts)
+ elif content is None:
+ _flat_result = ""
+ elif isinstance(content, str):
+ _flat_result = content
+ else:
+ _flat_result = _json.dumps(content)
+ filtered.append(
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": _tr_id,
+ "content": _flat_result,
+ }
+ ],
+ }
+ )
+ continue
if isinstance(content, list):
# Translate OpenAI multimodal parts -> Anthropic native shapes.
# - `image_url` -> `{type:"image", source:...}`
@@ -1583,6 +1913,37 @@ class ExternalProviderClient:
if title:
doc_block["title"] = title
anthropic_parts.append(doc_block)
+ # Assistant tool_calls -> Anthropic tool_use blocks
+ # appended to the same message. Anthropic native
+ # Messages API does not accept OpenAI's top-level
+ # `tool_calls` field; the call lives inside a content
+ # block with `{type:"tool_use", id, name, input}`.
+ if msg.get("role") == "assistant" and isinstance(
+ msg.get("tool_calls"), list
+ ):
+ for _tc in msg["tool_calls"]:
+ if not isinstance(_tc, dict):
+ continue
+ _fn = _tc.get("function") or {}
+ if not isinstance(_fn, dict) or not _fn.get("name"):
+ continue
+ _raw = _fn.get("arguments") or "{}"
+ try:
+ _input = (
+ _json.loads(_raw) if isinstance(_raw, str) else _raw
+ )
+ except Exception:
+ _input = {"_raw": _raw}
+ if not isinstance(_input, dict):
+ _input = {"value": _input}
+ anthropic_parts.append(
+ {
+ "type": "tool_use",
+ "id": _tc.get("id") or f"toolu_{time.time_ns()}",
+ "name": _fn["name"],
+ "input": _input,
+ }
+ )
# Skip whole-message append when nothing usable survived.
# An empty content array (e.g. user dropped only an unparseable
# `input_document`) would 400 the Anthropic API with
@@ -1590,6 +1951,72 @@ class ExternalProviderClient:
if anthropic_parts:
filtered.append({"role": msg["role"], "content": anthropic_parts})
else:
+ # role="tool" follow-up -> Anthropic native tool_result
+ # block on a `user` message. The OpenAI shape
+ # (role=tool, content=string, tool_call_id) is not a
+ # valid Anthropic role.
+ if msg.get("role") == "tool":
+ _tr_id = msg.get("tool_call_id") or ""
+ _tr_content = msg.get("content")
+ if _tr_content is None:
+ _tr_content = ""
+ filtered.append(
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "tool_result",
+ "tool_use_id": _tr_id,
+ "content": (
+ _tr_content
+ if isinstance(_tr_content, str)
+ else _json.dumps(_tr_content)
+ ),
+ }
+ ],
+ }
+ )
+ continue
+ # Assistant turn whose content is a plain string but
+ # also carries OpenAI `tool_calls`: convert into a
+ # content-array message with a text block + tool_use
+ # blocks. Without this, the top-level tool_calls leaks
+ # through unchanged.
+ if (
+ msg.get("role") == "assistant"
+ and isinstance(msg.get("tool_calls"), list)
+ and msg["tool_calls"]
+ ):
+ _text_content = msg.get("content")
+ _blocks: list[dict[str, Any]] = []
+ if isinstance(_text_content, str) and _text_content:
+ _blocks.append({"type": "text", "text": _text_content})
+ for _tc in msg["tool_calls"]:
+ if not isinstance(_tc, dict):
+ continue
+ _fn = _tc.get("function") or {}
+ if not isinstance(_fn, dict) or not _fn.get("name"):
+ continue
+ _raw = _fn.get("arguments") or "{}"
+ try:
+ _input = (
+ _json.loads(_raw) if isinstance(_raw, str) else _raw
+ )
+ except Exception:
+ _input = {"_raw": _raw}
+ if not isinstance(_input, dict):
+ _input = {"value": _input}
+ _blocks.append(
+ {
+ "type": "tool_use",
+ "id": _tc.get("id") or f"toolu_{time.time_ns()}",
+ "name": _fn["name"],
+ "input": _input,
+ }
+ )
+ if _blocks:
+ filtered.append({"role": "assistant", "content": _blocks})
+ continue
filtered.append(msg)
# Claude 4.7 family removed temperature / top_p / top_k entirely.
@@ -1616,34 +2043,16 @@ class ExternalProviderClient:
# same as True here (callers that don't set the flag still get
# caching). Pass False explicitly to opt out.
prompt_caching_enabled = enable_prompt_caching is not False
- # Anthropic accepts an optional `ttl` on each cache_control marker
- # (default is the 5m ephemeral pool; set "1h" to land in the 1h
- # pool instead). Per the prompt-caching docs, 1h cache writes are
- # billed at 2x base input vs 1.25x for 5m, but reads are 0.1x for
- # both. The 1h pool is the right pick when conversations span
- # multiple short bursts more than 5 minutes apart -- the read
- # discount makes up for the 1.6x write premium after a single
- # additional hit. Anything other than the known TTL strings is
- # dropped to avoid sending a malformed marker.
- #
- # The `extended-cache-ttl-2025-04-11` beta header that originally
- # gated 1h TTL has been promoted to GA: as of 2026-05 the live
- # API accepts `ttl: "1h"` without any beta opt-in. Verified
- # against api.anthropic.com on claude-opus-4-7 (status 200 +
- # `ephemeral_1h_input_tokens` populated). The test below pins
- # the contract by asserting the header is NOT on the wire so a
- # future regression that reintroduces the gate would surface
- # before users see a 400.
+ # Optional 1h cache TTL is GA as of 2026-05 (no beta header). 1h
+ # writes are 2x vs 5m's 1.25x but reads are 0.1x for both, so 1h
+ # wins after a single extra hit. Unknown TTL strings drop.
cache_marker: dict[str, Any] = {"type": "ephemeral"}
if prompt_cache_ttl in ("5m", "1h"):
cache_marker["ttl"] = prompt_cache_ttl
if system:
if prompt_caching_enabled:
- # System block is the most stable prefix across turns, so
- # it gets its own breakpoint. Skipped when system is
- # empty — there's nothing to cache, and an empty marker
- # is a no-op.
+ # System is the most stable cross-turn prefix; own breakpoint.
body["system"] = [
{
"type": "text",
@@ -1749,19 +2158,29 @@ class ExternalProviderClient:
if body.get("max_tokens", 0) <= budget_tokens:
body["max_tokens"] = budget_tokens + 1024
- # Anthropic server-side web_search — see
- # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool
- # The tool type is date-pinned per model family. Newer Opus /
- # Sonnet 4.6 + 4.7 accept `web_search_20260209` with dynamic
- # filtering (Claude writes code to filter results before they
- # reach context); everything else uses `web_search_20250305`.
- # `_anthropic_web_search_version` picks the right one. Anthropic
- # dispatches search calls server-side, returning server_tool_use
- # + web_search_tool_result blocks in the SSE stream, plus
- # url-citation annotations on text deltas. We translate all of
- # that into our local _toolEvent shape so the chat UI renders
- # web_search exactly like OpenAI's path.
- if enabled_tools and "web_search" in enabled_tools:
+ # tool_choice="none" or pinned-function suppresses hosted tools
+ # so a stale UI toggle can't fire server-side search/code-exec.
+ _anthropic_tool_choice_disabled = (
+ isinstance(tool_choice, str) and tool_choice.strip().lower() == "none"
+ )
+ _anthropic_tool_choice_forced_function = (
+ isinstance(tool_choice, dict)
+ and tool_choice.get("type") == "function"
+ and isinstance(tool_choice.get("function"), dict)
+ and bool(tool_choice["function"].get("name"))
+ )
+ _anthropic_hosted_builtins_allowed = (
+ not _anthropic_tool_choice_disabled
+ and not _anthropic_tool_choice_forced_function
+ )
+
+ # Anthropic web_search (date-pinned per model family).
+ # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool
+ if (
+ _anthropic_hosted_builtins_allowed
+ and enabled_tools
+ and "web_search" in enabled_tools
+ ):
anthropic_tools = list(body.get("tools") or [])
anthropic_tools.append(
{
@@ -1772,14 +2191,13 @@ class ExternalProviderClient:
)
body["tools"] = anthropic_tools
- # Anthropic server-side web_fetch reads a single URL (text/PDF)
- # and returns a `web_fetch_tool_result` document block. Opt in
- # via `enabled_tools=["web_fetch"]`; no beta header required.
- # `_anthropic_web_fetch_version` picks `web_fetch_20260209`
- # (dynamic filtering) for Opus 4.6/4.7 + Sonnet 4.6, falling
- # back to `web_fetch_20250910` elsewhere; mismatched variants
- # return 400 so the per-model picker is required.
- web_fetch_enabled = bool(enabled_tools and "web_fetch" in enabled_tools)
+ # Anthropic web_fetch: only URLs already in conversation. Date-pinned.
+ # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-fetch-tool
+ web_fetch_enabled = bool(
+ _anthropic_hosted_builtins_allowed
+ and enabled_tools
+ and "web_fetch" in enabled_tools
+ )
if web_fetch_enabled:
anthropic_tools = list(body.get("tools") or [])
anthropic_tools.append(
@@ -1792,24 +2210,13 @@ class ExternalProviderClient:
body["tools"] = anthropic_tools
# Anthropic server-side code execution — see
- # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool
- # The tool type is date-pinned per model family.
- # `_anthropic_code_execution_version` picks `code_execution_20260120`
- # for Opus 4.5+ / Sonnet 4.5+ / Opus 4.7 / Sonnet 4.6 (adds REPL
- # state persistence + programmatic tool calling) and falls back
- # to `code_execution_20250825` everywhere else. Both versions
- # run Python + bash + str_replace file edits inside a 5 GB
- # sandboxed container per request, with no internet access, and
- # both are unlocked by the same `code-execution-2025-08-25`
- # `anthropic-beta` header set further down. On the SSE stream
- # Anthropic emits two sub-tool names -- `bash_code_execution`
- # and `text_editor_code_execution` -- wrapped in the standard
- # server_tool_use / *_tool_result block shape.
- # v1 wires the tool only; file uploads (container_upload
- # content blocks and generated-file retrieval via the Files
- # API) are a deliberate follow-up.
+ # https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool
+ # Date-pinned tool type per model; both unlock via the same
+ # `code-execution-2025-08-25` beta header set below.
code_execution_enabled = bool(
- enabled_tools and "code_execution" in enabled_tools
+ _anthropic_hosted_builtins_allowed
+ and enabled_tools
+ and "code_execution" in enabled_tools
)
if code_execution_enabled:
anthropic_tools = list(body.get("tools") or [])
@@ -1820,32 +2227,14 @@ class ExternalProviderClient:
}
)
body["tools"] = anthropic_tools
- # Reuse the prior turn's container so filesystem state
- # (files written, packages installed, variables set)
- # persists across turns of the same thread. Anthropic
- # exposes the container id on the Message object's
- # top-level `container.id`; on the SSE stream we latch it
- # off `message_start.message.container.id` further down
- # and emit a `container_ready` _toolEvent so the chat
- # adapter persists it on the thread record. A stale id
- # (container expired / not found) surfaces as a 4xx
- # below, where we emit `container_invalidated` and let
- # the next turn fall back to auto-create.
+ # Reuse the thread's prior container so filesystem state
+ # persists. Stale ids 4xx and clear via container_invalidated.
if anthropic_code_exec_container_id:
body["container"] = anthropic_code_exec_container_id
- # Server-side context compaction — see
- # https://platform.claude.com/docs/en/build-with-claude/compaction
- # Beta as of `compact-2026-01-12`. When `compaction_threshold` is
- # provided AND the model accepts compaction (Opus 4.6+ / 4.7,
- # Sonnet 4.6, Mythos preview), attach
- # `context_management.edits[{type:"compact_20260112", trigger:
- # {type:"input_tokens", value:N}}]` to the body. Anthropic runs
- # the compaction step server-side once the rendered prompt
- # crosses the threshold and replies with a top-level
- # `context_management` block plus `usage.iterations[]` so we can
- # account per-iteration. Below-min thresholds get clamped up to
- # 50K so the request doesn't 400.
+ # Server-side compaction (beta `compact-2026-01-12`). Clamps
+ # below-min thresholds to 50K so the request doesn't 400.
+ # https://platform.claude.com/docs/en/build-with-claude/compaction
compaction_active = (
compaction_threshold is not None
and compaction_threshold > 0
@@ -1878,10 +2267,8 @@ class ExternalProviderClient:
url = f"{self.base_url}/messages"
completion_id = f"chatcmpl-anthropic-{model.replace('/', '-')}"
- # Log the outgoing config keys (not the messages themselves) so we
- # can prove which thinking/effort fields actually reached the wire.
- # If Anthropic skips reasoning despite a configured effort, this
- # tells us whether we sent the field or dropped it on the floor.
+ # Log outgoing config keys (not messages) to prove which thinking /
+ # effort fields actually reached the wire.
logger.info(
"Anthropic request shape (model=%s, has_thinking=%s, thinking=%s, "
"output_config=%s, temperature=%s, has_top_p=%s, has_top_k=%s, "
@@ -1896,16 +2283,10 @@ class ExternalProviderClient:
body.get("max_tokens"),
)
- # Translate Anthropic stop reasons onto the OpenAI chat-completions
- # `finish_reason` vocabulary. `pause_turn` maps to None so the
- # adapter does NOT emit a finish_reason chunk: pause_turn means
- # Claude paused a long server-tool turn (web_search / web_fetch)
- # and will continue once the user (or our retry) sends back the
- # partial assistant message. Forwarding it as "stop" makes the
- # OpenAI client think the answer is done and truncates the
- # rendered message. `refusal` maps to "content_filter" as the
- # nearest semantic match. See
- # https://platform.claude.com/docs/en/api/messages#response-stop-reason
+ # Anthropic stop_reason -> OpenAI finish_reason. `pause_turn`
+ # maps to None so the UI doesn't treat a paused server-tool turn
+ # as final. `refusal` -> "content_filter" (closest match).
+ # https://platform.claude.com/docs/en/api/messages#response-stop-reason
_finish_reason_map: dict[str, Optional[str]] = {
"end_turn": "stop",
"max_tokens": "length",
@@ -1918,11 +2299,7 @@ class ExternalProviderClient:
logger.info("Proxying Anthropic Messages API to %s (model=%s)", url, model)
request_headers = self._auth_headers()
- # Anthropic accepts comma-separated beta features in a single
- # `anthropic-beta` header. Merge our flags onto whatever the
- # registry's extra_headers contributed (currently nothing on
- # the beta axis, just anthropic-version) so future betas
- # added at the registry level keep working.
+ # Merge new beta flags onto whatever the registry contributed.
existing_beta = request_headers.get("anthropic-beta", "").strip()
beta_parts = (
[p.strip() for p in existing_beta.split(",") if p.strip()]
@@ -1988,27 +2365,14 @@ class ExternalProviderClient:
# "no thinking content" — distinguishes "Anthropic never sent
# thinking_delta" from "frontend didn't render the chunks".
event_counts: dict[str, int] = {}
- # web_search state. Anthropic emits the query inside an
- # `input_json_delta` stream on a `server_tool_use` content
- # block, then a separate `web_search_tool_result` block
- # with the URL list. Unlike OpenAI we get per-call results
- # directly, so each tool card carries its own citations.
- # `current_server_tool_use`: {id, name, partial_json_buffer}
- # `current_result_block`: {tool_use_id, results}
- # Both go to None when the matching content_block_stop fires.
+ # web_search state. Query streams via input_json_delta
+ # on a server_tool_use block; results land in a separate
+ # web_search_tool_result block. Per-call citations.
current_server_tool_use: Optional[dict[str, Any]] = None
current_result_block: Optional[dict[str, Any]] = None
web_search_calls: dict[str, dict[str, Any]] = {}
- # code_execution state. Anthropic's
- # `code_execution_20250825` tool emits the same
- # server_tool_use → *_tool_result block shape as
- # web_search, but the server_tool_use carries one of
- # two sub-tool names (`bash_code_execution` or
- # `text_editor_code_execution`) and the result block
- # type matches (`bash_code_execution_tool_result` /
- # `text_editor_code_execution_tool_result`). Kept
- # parallel to web_search state so the two paths don't
- # collide when both pills are on in the same turn.
+ # code_execution state (bash / text_editor sub-tools);
+ # kept parallel to web_search so concurrent pills don't collide.
current_code_exec_use: Optional[dict[str, Any]] = None
current_code_exec_result: Optional[dict[str, Any]] = None
code_execution_calls: dict[str, dict[str, Any]] = {}
@@ -2078,6 +2442,7 @@ class ExternalProviderClient:
return f"data: {_json.dumps(chunk)}"
def _emit_tool_event(payload: dict[str, Any]) -> str:
+ _stamp_server_tool_marker(payload)
chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
@@ -2335,18 +2700,8 @@ class ExternalProviderClient:
"inner": inner if isinstance(inner, dict) else {},
}
elif block_type == "compaction":
- # Server-side compaction emits a `compaction`
- # content block on the assistant message.
- # Anthropic may include the summary text on
- # this start event AND/OR stream it via
- # text_delta events on the same block. See
- # https://platform.claude.com/docs/en/build-with-claude/compaction
- # Capture either form; finalize and emit
- # on content_block_stop. The chat-adapter
- # persists the block onto the assistant
- # message so the next turn's request
- # carries it back -- Anthropic then skips
- # re-compaction from scratch.
+ # Summary may arrive on start AND/OR via
+ # text_delta. Capture both; emit on stop.
seed = content_block.get("content") or ""
current_compaction = {
"content": seed if isinstance(seed, str) else "",
@@ -2356,12 +2711,7 @@ class ExternalProviderClient:
delta = event.get("delta", {})
delta_type = delta.get("type")
if delta_type == "thinking_delta":
- # Anthropic streams extended-thinking content as
- # thinking_delta events on a separate content
- # block. Wrap as inline ... so
- # the frontend's parseAssistantContent lifts it
- # into the reasoning panel — same pattern as
- # the OpenAI Responses path.
+ # Wrap as ... for parseAssistantContent.
thinking_text = delta.get("thinking", "")
if thinking_text:
if not thinking_open:
@@ -2631,19 +2981,9 @@ class ExternalProviderClient:
delta_usage = event.get("usage")
if isinstance(delta_usage, dict):
last_usage.update(delta_usage)
- # When a fresh compaction has run, Anthropic
- # publishes per-iteration token counts in
- # `usage.iterations[]`. The top-level
- # input_tokens / output_tokens only cover the
- # `message` iteration, NOT the compaction
- # passes — billing has to sum the whole
- # array. See
- # https://platform.claude.com/docs/en/build-with-claude/compaction
- # Fold the compaction iterations into
- # `compaction_input_tokens` / `compaction_output_tokens`
- # so the cost surface can add them without
- # re-walking the array (and so the closing
- # log line names the figures).
+ # Compaction iterations aren't in top-level
+ # input/output_tokens; fold them into
+ # compaction_{input,output}_tokens for billing.
iterations = delta_usage.get("iterations")
if isinstance(iterations, list):
c_in = 0
@@ -2880,6 +3220,1722 @@ class ExternalProviderClient:
self.provider_type,
)
+ async def _stream_gemini(
+ self,
+ messages: list[dict[str, Any]],
+ model: str,
+ temperature: float,
+ top_p: float,
+ max_tokens: Optional[int],
+ top_k: Optional[int] = None,
+ presence_penalty: float = 0.0,
+ enabled_tools: Optional[list[str]] = None,
+ enable_prompt_caching: Optional[Any] = None,
+ enable_thinking: Optional[bool] = None,
+ reasoning_effort: Optional[str] = None,
+ tools: Optional[list[dict[str, Any]]] = None,
+ tool_choice: Optional[Any] = None,
+ ) -> AsyncGenerator[str, None]:
+ """
+ Call Google's native Gemini API and translate its streaming
+ ``streamGenerateContent`` response into OpenAI Chat Completions
+ chunk format.
+
+ Gemini does NOT speak the OpenAI Chat Completions contract on
+ its primary endpoint. The wire shape is:
+
+ POST /v1beta/models/{model}:streamGenerateContent?alt=sse
+ {
+ "contents": [{"role": "user|model", "parts": [{"text": "..."}]}],
+ "systemInstruction": {"parts": [{"text": "..."}]},
+ "generationConfig": {"temperature": 0.7, "topP": 0.95, "topK": 40,
+ "maxOutputTokens": 1024},
+ "tools": [{"googleSearch": {}}, {"codeExecution": {}}],
+ "cachedContent": "" // optional, see caching docs
+ }
+
+ Streamed responses are SSE frames carrying partial
+ ``GenerateContentResponse`` objects:
+
+ {"candidates": [{"content": {"parts": [{"text": "Hello"}]},
+ "finishReason": "STOP"}],
+ "usageMetadata": {"promptTokenCount": 7, "candidatesTokenCount": 3}}
+
+ Image generation uses the same endpoint with model
+ ``gemini-2.5-flash-image`` (also called Nano Banana); the
+ response carries an ``inlineData`` part with the base64 PNG
+ bytes and a ``mimeType``. We surface that through the same
+ ``tool_start`` / ``tool_end`` ``image_b64`` envelope the OpenAI
+ image_generation path uses, so the chat UI renders the image
+ inline with no extra plumbing.
+
+ References:
+ - https://ai.google.dev/gemini-api/docs/text-generation
+ - https://ai.google.dev/gemini-api/docs/function-calling
+ - https://ai.google.dev/gemini-api/docs/grounding
+ - https://ai.google.dev/gemini-api/docs/caching
+ - https://ai.google.dev/gemini-api/docs/image-generation
+ """
+ import json as _json
+
+ # Validate the user-controlled model id BEFORE any message
+ # translation. A model like `../cachedContents/x` is path-
+ # traversal that lands in `/v1beta/cachedContents/...`; rejecting
+ # it here also avoids triggering user-controlled outbound fetches
+ # (remote image_url inlining) on a request we'll error out
+ # anyway. Documented catalog ids match `[A-Za-z0-9._-]+`.
+ if not re.fullmatch(r"[A-Za-z0-9._-]+", model):
+ yield _error_sse_line(
+ 400,
+ f"Invalid Gemini model id: {model!r}",
+ self.provider_type,
+ )
+ return
+
+ # Translate OpenAI messages -> Gemini contents. system role
+ # promotes to top-level systemInstruction.
+ system_text_parts: list[str] = []
+ contents: list[dict[str, Any]] = []
+ # OpenAI may drop `name` from role="tool" follow-ups. Remember
+ # prior function names so functionResponse isn't sent name-less
+ # (Gemini 400s on empty names).
+ tool_call_names: dict[str, str] = {}
+ # tool_call_ids whose assistant card was dropped (synthetic
+ # builtin) or already replayed as native parts. Their role="tool"
+ # follow-up must be skipped to avoid orphan/duplicate responses.
+ _gemini_skip_tool_result_ids: set[str] = set()
+ # Per-request image caps. The byte cap counts DECODED bytes; we
+ # set it to ~14 MB because base64 expansion + prompt overhead
+ # must fit Gemini's ~20 MB request limit.
+ _GEMINI_REMOTE_IMAGE_MAX_COUNT = 8
+ _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES = 14 * 1024 * 1024
+ _remote_image_count = 0
+ _remote_image_total_bytes = 0
+ for msg in messages:
+ role = msg.get("role")
+ content = msg.get("content", "")
+ if role == "system":
+ if isinstance(content, str):
+ if content:
+ system_text_parts.append(content)
+ elif isinstance(content, list):
+ for part in content:
+ if (
+ isinstance(part, dict)
+ and part.get("type") == "text"
+ and part.get("text")
+ ):
+ system_text_parts.append(part["text"])
+ continue
+ # Map OpenAI roles to Gemini's two-role contract.
+ gemini_role = "model" if role == "assistant" else "user"
+ parts: list[dict[str, Any]] = []
+ if isinstance(content, str):
+ if content:
+ parts.append({"text": content})
+ elif isinstance(content, list):
+ for part in content:
+ if not isinstance(part, dict):
+ continue
+ ptype = part.get("type")
+ if ptype == "text":
+ text = part.get("text", "")
+ if text:
+ parts.append({"text": text})
+ elif ptype == "image_url":
+ url = part.get("image_url", {}).get("url", "")
+ if url.startswith("data:"):
+ header, _, b64data = url.partition(",")
+ media_type = (
+ header.split(";")[0]
+ .replace("data:", "")
+ .strip()
+ .lower()
+ or "image/jpeg"
+ )
+ # Symmetry with the fetched remote image
+ # path, which already rejects non-image
+ # Content-Type. A `data:text/html;base64,...`
+ # URL otherwise lands as Gemini inlineData
+ # with mimeType="text/html" and 400s the
+ # whole request.
+ if not media_type.startswith("image/"):
+ logger.info(
+ "Gemini inlineData: refusing non-image data URL media_type=%s",
+ media_type,
+ )
+ elif b64data:
+ # data: URLs share the same caps as fetched
+ # URLs so inline payloads don't bypass them.
+ _data_approx_bytes = (len(b64data) * 3) // 4
+ if (
+ _remote_image_count
+ >= _GEMINI_REMOTE_IMAGE_MAX_COUNT
+ ):
+ logger.info(
+ "Gemini inlineData: per-request count cap %d reached, dropping image",
+ _GEMINI_REMOTE_IMAGE_MAX_COUNT,
+ )
+ elif (
+ _remote_image_total_bytes + _data_approx_bytes
+ > _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES
+ ):
+ logger.info(
+ "Gemini inlineData: per-request byte cap reached, dropping image",
+ )
+ else:
+ _remote_image_count += 1
+ _remote_image_total_bytes += _data_approx_bytes
+ parts.append(
+ {
+ "inlineData": {
+ "mimeType": media_type,
+ "data": b64data,
+ }
+ }
+ )
+ elif url:
+ # fileData.fileUri only accepts Files-API URIs
+ # and YouTube; everything else must be downloaded
+ # and inlined. Parse fields explicitly so
+ # attacker URLs like https://evil.com/youtube.com/x
+ # aren't misclassified as YouTube.
+ try:
+ _parsed_image_url = urlparse(url)
+ except (ValueError, UnicodeError):
+ _parsed_image_url = None
+ if _parsed_image_url is None:
+ _img_scheme = ""
+ _img_host = ""
+ _img_path = ""
+ else:
+ _img_scheme = (_parsed_image_url.scheme or "").lower()
+ _img_host = (_parsed_image_url.hostname or "").lower()
+ _img_path = _parsed_image_url.path or ""
+ _is_native_uri = (
+ _img_scheme == "https"
+ and _img_host == "generativelanguage.googleapis.com"
+ and _img_path.startswith("/v1beta/files/")
+ )
+ _is_youtube = _img_scheme == "https" and (
+ _img_host == "youtu.be"
+ or _img_host == "youtube.com"
+ or _img_host.endswith(".youtube.com")
+ )
+ _guessed, _ = mimetypes.guess_type(_img_path)
+ _media_type = (
+ _guessed
+ if isinstance(_guessed, str)
+ and _guessed.startswith("image/")
+ else "image/jpeg"
+ )
+ if _is_youtube:
+ # YouTube URIs must use video/mp4; the
+ # default image/jpeg yields a 400.
+ parts.append(
+ {
+ "fileData": {
+ "fileUri": url,
+ "mimeType": "video/mp4",
+ }
+ }
+ )
+ elif _is_native_uri:
+ parts.append(
+ {
+ "fileData": {
+ "fileUri": url,
+ "mimeType": _media_type,
+ }
+ }
+ )
+ elif _remote_image_count >= _GEMINI_REMOTE_IMAGE_MAX_COUNT:
+ logger.info(
+ "Gemini image fetch: per-request count cap %d reached, dropping image",
+ _GEMINI_REMOTE_IMAGE_MAX_COUNT,
+ )
+ else:
+ # Refuse pre-fetch when the per-request
+ # byte budget is spent; pass the remainder
+ # so over-budget URLs reject on Content-Length.
+ _remaining_bytes = (
+ _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES
+ - _remote_image_total_bytes
+ )
+ if _remaining_bytes <= 0:
+ logger.info(
+ "Gemini image fetch: per-request byte cap already reached, dropping image",
+ )
+ else:
+ # Count attempts before awaiting so
+ # slow URLs don't each burn the timeout.
+ _remote_image_count += 1
+ _fetched = await _safe_fetch_image_for_gemini(
+ url,
+ _media_type,
+ max_bytes = _remaining_bytes,
+ )
+ if _fetched is not None:
+ _final_mime, _b64 = _fetched
+ # base64 expands ~4/3 — recover bytes from len(_b64).
+ _approx_bytes = (len(_b64) * 3) // 4
+ if (
+ _remote_image_total_bytes + _approx_bytes
+ > _GEMINI_REMOTE_IMAGE_MAX_TOTAL_BYTES
+ ):
+ logger.info(
+ "Gemini image fetch: per-request byte cap reached, dropping image",
+ )
+ else:
+ _remote_image_total_bytes += _approx_bytes
+ parts.append(
+ {
+ "inlineData": {
+ "mimeType": _final_mime,
+ "data": _b64,
+ }
+ }
+ )
+ # Gemini 3 strict function-calling requires text-part
+ # thoughtSignatures to be replayed on history; the frontend
+ # stows the latest one as
+ # extra_content.google.thought_signature on the assistant
+ # message and we pin it onto the last text part here.
+ if role == "assistant" and parts:
+ _msg_extra = msg.get("extra_content") if isinstance(msg, dict) else None
+ if isinstance(_msg_extra, dict):
+ _msg_g = _msg_extra.get("google") or {}
+ if isinstance(_msg_g, dict):
+ _msg_sig = _msg_g.get("thought_signature") or _msg_g.get(
+ "thoughtSignature"
+ )
+ if isinstance(_msg_sig, str) and _msg_sig:
+ for _idx in range(len(parts) - 1, -1, -1):
+ if "text" in parts[_idx]:
+ parts[_idx] = {
+ **parts[_idx],
+ "thoughtSignature": _msg_sig,
+ }
+ break
+ # Translate OpenAI tool_calls into Gemini functionCall parts.
+ # code_execution / image_generation replay their native parts
+ # (executableCode / codeExecutionResult / inlineData) stowed
+ # on extra_content.google.native_part.
+ tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
+ if isinstance(tool_calls, list):
+ for tc in tool_calls:
+ if not isinstance(tc, dict):
+ continue
+ fn = tc.get("function") or {}
+ if not isinstance(fn, dict):
+ continue
+ args_raw = fn.get("arguments") or "{}"
+ if isinstance(args_raw, str):
+ try:
+ args = _json.loads(args_raw)
+ except Exception:
+ args = {"_raw": args_raw}
+ elif isinstance(args_raw, dict):
+ args = args_raw
+ else:
+ args = {}
+ fn_name = fn.get("name", "")
+ tc_id = tc.get("id")
+ if fn_name and isinstance(tc_id, str) and tc_id:
+ tool_call_names[tc_id] = fn_name
+
+ # Replay native Gemini code_execution / image_generation parts
+ # from extra_content.google.native_part, with fallback to
+ # args.google.native_part for OAI-compat round-trips.
+ _extra = tc.get("extra_content")
+ _native_part = None
+ _google_extra: dict[str, Any] = {}
+ if isinstance(_extra, dict):
+ _ge = _extra.get("google") or {}
+ if isinstance(_ge, dict):
+ _google_extra = _ge
+ _native_part = _ge.get("native_part")
+ if _native_part is None and isinstance(args, dict):
+ _args_google = args.get("google")
+ if isinstance(_args_google, dict):
+ _args_np = _args_google.get("native_part")
+ if isinstance(_args_np, dict):
+ _native_part = _args_np
+ if not _google_extra:
+ _google_extra = _args_google
+
+ # Synthetic builtin cards (web_search/web_fetch) must
+ # not become fake functionCalls; drop them. Native
+ # code_execution / image_generation replay below.
+ _name_lc = fn_name.lower() if isinstance(fn_name, str) else ""
+ _is_synthetic_server_builtin = (
+ _name_lc
+ in (
+ "web_search",
+ "web_fetch",
+ "code_execution",
+ "image_generation",
+ )
+ and isinstance(args, dict)
+ and (
+ args.get("_server_tool") is True
+ or isinstance(
+ (args.get("google") or {}).get("native_part"), dict
+ )
+ )
+ )
+ if _is_synthetic_server_builtin and not (
+ _name_lc in ("code_execution", "image_generation")
+ and isinstance(_native_part, dict)
+ ):
+ # No replayable Gemini native part -- skip
+ # entirely rather than send a fake functionCall.
+ # Also remember this tool_call_id so a matching
+ # role="tool" follow-up does not become an
+ # orphan functionResponse below.
+ if isinstance(tc_id, str) and tc_id:
+ _gemini_skip_tool_result_ids.add(tc_id)
+ tool_call_names.pop(tc_id, None)
+ continue
+ if fn_name in ("code_execution", "image_generation") and isinstance(
+ _native_part, dict
+ ):
+ # code_execution/image_generation history is
+ # replayed as native parts; the matching
+ # role="tool" must be skipped or Gemini sees a
+ # functionResponse with no declared function
+ # name and 400s the turn.
+ if isinstance(tc_id, str) and tc_id:
+ _gemini_skip_tool_result_ids.add(tc_id)
+ # New shape: `native_part.parts` is an ordered list
+ # of full part wrappers, each carrying its own
+ # `thoughtSignature`. This preserves Gemini 3's
+ # strict per-part replay requirement when the
+ # frontend has merged executableCode +
+ # codeExecutionResult + inlineData into the same
+ # tool-call card.
+ _native_parts_list = _native_part.get("parts")
+ if isinstance(_native_parts_list, list):
+ for _entry in _native_parts_list:
+ if isinstance(_entry, dict):
+ parts.append(_entry)
+ continue
+ # Legacy single-object native_part: fan the shared
+ # thoughtSignature only when one subpart exists;
+ # for code+result, prefer executableCode and drop
+ # the signature elsewhere.
+ _legacy_sig = _native_part.get(
+ "thoughtSignature"
+ ) or _native_part.get("thought_signature")
+ _legacy_subparts = [
+ _k
+ for _k in (
+ "executableCode",
+ "codeExecutionResult",
+ "inlineData",
+ )
+ if isinstance(_native_part.get(_k), dict)
+ ]
+ for _native_key in (
+ "executableCode",
+ "codeExecutionResult",
+ "inlineData",
+ ):
+ _sub = _native_part.get(_native_key)
+ if not isinstance(_sub, dict):
+ continue
+ _replay_part: dict[str, Any] = {_native_key: _sub}
+ if isinstance(_legacy_sig, str) and _legacy_sig:
+ if len(_legacy_subparts) == 1:
+ _replay_part["thoughtSignature"] = _legacy_sig
+ elif _native_key == "executableCode":
+ _replay_part["thoughtSignature"] = _legacy_sig
+ parts.append(_replay_part)
+ continue
+
+ # Forward the OpenAI tool_call id into Gemini's
+ # functionCall.id so a follow-up turn that issues
+ # multiple calls to the same function (different
+ # args, same name) can be disambiguated on the
+ # response side. Gemini accepts the field per
+ # https://ai.google.dev/gemini-api/docs/function-calling.
+ function_call_part: dict[str, Any] = {
+ "name": fn_name,
+ "args": args,
+ }
+ if isinstance(tc_id, str) and tc_id:
+ function_call_part["id"] = tc_id
+ # Gemini 3 function-calling requires the prior
+ # thoughtSignature to be echoed back as a sibling
+ # of the functionCall part. The translator stows
+ # it on the assistant tool_call via
+ # `extra_content.google.thought_signature` (see
+ # the inbound emit below).
+ fc_part: dict[str, Any] = {"functionCall": function_call_part}
+ sig = _google_extra.get("thought_signature") or _google_extra.get(
+ "thoughtSignature"
+ )
+ if isinstance(sig, str) and sig:
+ fc_part["thoughtSignature"] = sig
+ parts.append(fc_part)
+ if role == "tool":
+ # If the matching assistant-side tool_call was either
+ # dropped (synthetic server-tool with no native part)
+ # or already replayed as Gemini-native parts
+ # (code_execution/image_generation native_part), drop
+ # the follow-up too. Emitting it as a functionResponse
+ # would be orphaned or duplicate the native result.
+ _tc_id_for_skip = msg.get("tool_call_id")
+ if (
+ isinstance(_tc_id_for_skip, str)
+ and _tc_id_for_skip in _gemini_skip_tool_result_ids
+ ):
+ continue
+ # OpenAI's role="tool" follow-up carries the function
+ # result. Gemini's matching shape is a role="user" turn
+ # with a functionResponse part. When the caller dropped
+ # ``name``, recover it from the matching assistant
+ # tool_call so Gemini doesn't 400 on an empty name.
+ tool_name = msg.get("name") or msg.get("tool_name") or ""
+ if not tool_name:
+ tc_id = msg.get("tool_call_id")
+ if isinstance(tc_id, str) and tc_id in tool_call_names:
+ tool_name = tool_call_names[tc_id]
+ response_payload: Any
+ if isinstance(content, list):
+ # OpenAI tool messages may carry list-form content
+ # (`[{"type":"text","text":"..."}]`). Forwarding the
+ # content-part objects verbatim into Gemini's
+ # `functionResponse.response.result` yields
+ # `result:[{"type":"text","text":"..."}]` instead of
+ # the actual tool output text; flatten text parts so
+ # the result mirrors the string-content path.
+ _flat_parts: list[str] = []
+ for _cpart in content:
+ if (
+ isinstance(_cpart, dict)
+ and _cpart.get("type") == "text"
+ and isinstance(_cpart.get("text"), str)
+ ):
+ _flat_parts.append(_cpart["text"])
+ _flat_text = "".join(_flat_parts)
+ try:
+ response_payload = _json.loads(_flat_text)
+ except Exception:
+ response_payload = {"result": _flat_text}
+ elif isinstance(content, str):
+ try:
+ response_payload = _json.loads(content)
+ except Exception:
+ response_payload = {"result": content}
+ else:
+ response_payload = content or {}
+ function_response_part: dict[str, Any] = {
+ "name": tool_name,
+ "response": (
+ response_payload
+ if isinstance(response_payload, dict)
+ else {"result": response_payload}
+ ),
+ }
+ # Mirror tool_call_id onto functionResponse.id so
+ # Gemini can match the result to the originating
+ # functionCall when multiple parallel calls were made.
+ tc_id = msg.get("tool_call_id")
+ if isinstance(tc_id, str) and tc_id:
+ function_response_part["id"] = tc_id
+ parts = [{"functionResponse": function_response_part}]
+ gemini_role = "user"
+ if parts:
+ # Gemini expects parallel functionResponses (multiple
+ # OpenAI role="tool" messages in a row) to ride on a
+ # single user content with multiple functionResponse
+ # parts -- the docs show parallel responses grouped
+ # together in the next turn. Merge consecutive
+ # functionResponse-only user blocks so realistic
+ # parallel tool loops round-trip correctly.
+ if (
+ role == "tool"
+ and contents
+ and contents[-1].get("role") == "user"
+ and all(
+ isinstance(p, dict) and "functionResponse" in p
+ for p in (contents[-1].get("parts") or [])
+ )
+ ):
+ contents[-1]["parts"].extend(parts)
+ else:
+ contents.append({"role": gemini_role, "parts": parts})
+
+ body: dict[str, Any] = {"contents": contents}
+ if system_text_parts:
+ body["systemInstruction"] = {
+ "parts": [{"text": "\n\n".join(system_text_parts)}]
+ }
+
+ # Generation config -- temperature / topP / topK / maxOutputTokens
+ # map straight across. The frontend capability matrix restricts
+ # the sliders the UI exposes for Gemini to this set.
+ gen_config: dict[str, Any] = {}
+ if temperature is not None:
+ gen_config["temperature"] = temperature
+ if top_p is not None:
+ gen_config["topP"] = top_p
+ if top_k is not None and top_k > 0:
+ gen_config["topK"] = top_k
+ # Gemini accepts ``presencePenalty`` on generationConfig with the
+ # same sign convention as the OpenAI knob (positive discourages
+ # repetition). Forward when the caller bothers to set it.
+ if presence_penalty:
+ gen_config["presencePenalty"] = presence_penalty
+ if max_tokens is not None:
+ gen_config["maxOutputTokens"] = max_tokens
+
+ # Nano Banana image generation. Gemini only accepts
+ # `responseModalities: ["TEXT","IMAGE"]` on the image-capable
+ # model family (id contains `-image` or `nano-banana`). Text-
+ # only models such as `gemini-2.5-flash` 400 on the same body,
+ # so only force image mode when the selected model actually
+ # supports it -- a stale `enabled_tools=["image_generation"]`
+ # on a text model is silently treated as a regular turn.
+ # https://ai.google.dev/gemini-api/docs/image-generation
+ model_lc = model.lower()
+ is_image_picker_model = "-image" in model_lc or "nano-banana" in model_lc
+ # tool_choice="none" / forced-function tool_choice must also
+ # suppress the implicit image-generation hosted tool. Otherwise
+ # an explicit OpenAI-style opt-out (or an explicit user-function
+ # pin) still flips `responseModalities=["TEXT","IMAGE"]` on
+ # image-tier models and bills for image output.
+ _tool_choice_disabled = (
+ isinstance(tool_choice, str) and tool_choice.strip().lower() == "none"
+ )
+ _tool_choice_forced_function = (
+ isinstance(tool_choice, dict)
+ and tool_choice.get("type") == "function"
+ and isinstance(tool_choice.get("function"), dict)
+ and bool(tool_choice["function"].get("name"))
+ )
+ _hosted_builtins_allowed = (
+ not _tool_choice_disabled and not _tool_choice_forced_function
+ )
+ # Image-tier model IDs reject text-only tools (code_execution,
+ # user functions) and thinkingConfig regardless of whether the
+ # Images pill is on -- those are model-level constraints
+ # documented by Google. The pill only controls whether we ask
+ # Gemini to actually emit image output via
+ # `responseModalities: ["TEXT","IMAGE"]`. Decoupling the two
+ # avoids the case where Images is off + Code/Search is on
+ # forwards `tools: [{codeExecution: {}}]` plus
+ # `thinkingConfig` to an image model and 400s.
+ image_tool_requested = bool(
+ _hosted_builtins_allowed
+ and enabled_tools
+ and "image_generation" in enabled_tools
+ )
+ # Strict tool / thinking strip uses the model-id check.
+ is_image_model_strict = is_image_picker_model
+ # The actual modality flip only happens when the user opted in.
+ is_image_model = is_image_picker_model and image_tool_requested
+ if is_image_model:
+ gen_config["responseModalities"] = ["TEXT", "IMAGE"]
+ elif is_image_picker_model:
+ # Force TEXT-only so an image-capable model with Images OFF
+ # doesn't still bill for image output.
+ gen_config["responseModalities"] = ["TEXT"]
+
+ # Thinking control. Gemini 3 uses thinkingLevel (str), 2.5 uses
+ # thinkingBudget (int). Gemini 3 has no full-off; minimum is
+ # "minimal" on Flash, "low" on Pro.
+ # https://ai.google.dev/gemini-api/docs/thinking
+ _GEMINI3_THINKING_PREFIXES = (
+ "gemini-3.5-",
+ "gemini-3.1-",
+ "gemini-3-",
+ "gemini-pro-latest",
+ "gemini-flash-latest",
+ "gemini-flash-lite-latest",
+ )
+ _GEMINI3_PRO_PREFIXES = (
+ "gemini-3.5-pro",
+ "gemini-3.1-pro",
+ "gemini-3-pro",
+ "gemini-pro-latest",
+ )
+ _PRO_THINKING_PREFIXES = ("gemini-2.5-pro",)
+ is_gemini3_thinking = any(
+ model_lc.startswith(p) for p in _GEMINI3_THINKING_PREFIXES
+ )
+ is_gemini3_pro = any(model_lc.startswith(p) for p in _GEMINI3_PRO_PREFIXES)
+ _is_pro_thinking_only = any(
+ model_lc == p or model_lc.startswith(p + "-")
+ for p in _PRO_THINKING_PREFIXES
+ )
+ effort_lc = (reasoning_effort or "").strip().lower()
+ if not is_image_model_strict and is_gemini3_thinking:
+ # Gemini 3.x thinkingLevel matrix:
+ # 3.1+ Pro: low/medium/high
+ # 3 Pro: low/high (deprecated 2026-03-09)
+ # 3.x Flash*: minimal/low/medium/high
+ # Coerce minimal->low on Pro; medium->high on legacy 3-Pro.
+ _G3_LEVELS = {"minimal", "low", "medium", "high"}
+ level: Optional[str] = None
+ if effort_lc in ("none", "off"):
+ level = "low" if is_gemini3_pro else "minimal"
+ elif effort_lc == "max":
+ level = "high"
+ elif effort_lc in _G3_LEVELS:
+ # Coerce legacy 3-Pro (low/high only) inputs.
+ _is_legacy_gemini3_pro = model_lc.startswith(
+ ("gemini-3-pro-preview", "gemini-3-pro")
+ ) and not model_lc.startswith(("gemini-3.1-pro", "gemini-3.5-pro"))
+ if is_gemini3_pro and effort_lc == "minimal":
+ level = "low"
+ elif _is_legacy_gemini3_pro and effort_lc == "medium":
+ level = "high"
+ else:
+ level = effort_lc
+ elif enable_thinking is True:
+ level = "high"
+ elif enable_thinking is False:
+ level = "low" if is_gemini3_pro else "minimal"
+ if level is not None:
+ gen_config["thinkingConfig"] = {"thinkingLevel": level}
+ elif not is_image_model_strict:
+ # Gemini 2.5 / older: thinkingBudget int. Effort -> budget
+ # mirrors the OpenAI minimal/low/medium/high ladder so the
+ # existing frontend picker maps cleanly.
+ # NOTE: gemini-2.5-flash-lite rejects positive budgets below
+ # 512 with HTTP 400, so minimal=512 sits at that floor.
+ _EFFORT_TO_BUDGET: dict[str, int] = {
+ "minimal": 512,
+ "low": 2048,
+ "medium": 8192,
+ "high": 24576,
+ "xhigh": -1,
+ "max": -1,
+ }
+ thinking_budget: Optional[int] = None
+ if effort_lc == "none" or enable_thinking is False:
+ # Pro-tier 2.5 rejects budget=0 (400 "only works in
+ # thinking mode"), so coerce to a small positive value.
+ thinking_budget = 128 if _is_pro_thinking_only else 0
+ elif effort_lc in _EFFORT_TO_BUDGET:
+ thinking_budget = _EFFORT_TO_BUDGET[effort_lc]
+ elif enable_thinking is True:
+ thinking_budget = -1
+ if thinking_budget is not None:
+ gen_config["thinkingConfig"] = {
+ "thinkingBudget": thinking_budget,
+ }
+
+ if gen_config:
+ body["generationConfig"] = gen_config
+
+ # Hosted tools: googleSearch (grounding) and codeExecution.
+ # Image-mode rejects codeExecution; only Gemini 3 image models
+ # accept googleSearch.
+ # https://ai.google.dev/gemini-api/docs/grounding
+ # https://ai.google.dev/gemini-api/docs/code-execution
+ def _gemini_image_model_allows_google_search(_m: str) -> bool:
+ return (
+ _m.startswith("gemini-3-pro-image")
+ or _m.startswith("gemini-3.1-flash-image")
+ or _m.startswith("nano-banana-pro")
+ or _m.startswith("nano-banana-2")
+ )
+
+ google_search_allowed = (
+ not is_image_model_strict
+ or _gemini_image_model_allows_google_search(model_lc)
+ )
+ code_execution_allowed = not is_image_model_strict
+ text_tools_allowed = not is_image_model_strict
+ # tool_choice="none" / forced-function suppresses hosted builtins
+ # too, matching the Anthropic / OpenRouter gates.
+ tools_array: list[dict[str, Any]] = []
+ if (
+ _hosted_builtins_allowed
+ and enabled_tools
+ and "web_search" in enabled_tools
+ and google_search_allowed
+ ):
+ tools_array.append({"googleSearch": {}})
+ if (
+ _hosted_builtins_allowed
+ and enabled_tools
+ and "code_execution" in enabled_tools
+ and code_execution_allowed
+ ):
+ tools_array.append({"codeExecution": {}})
+ # OpenAI-style function declarations -> Gemini functionDeclarations.
+ # https://ai.google.dev/gemini-api/docs/function-calling#step_1
+ # Gemini's Schema accepts only the OpenAPI 3.0 subset documented
+ # at https://ai.google.dev/api/caching#Schema; OpenAI's strict
+ # tool definitions routinely include `additionalProperties`,
+ # `$schema`, `$defs`, `strict`, `examples`, and similar keys
+ # which 400 the request as INVALID_ARGUMENT. Strip them
+ # recursively before forwarding.
+ _GEMINI_ALLOWED_SCHEMA_KEYS = frozenset(
+ {
+ "type",
+ "format",
+ "title",
+ "description",
+ "nullable",
+ "enum",
+ "maxItems",
+ "minItems",
+ "properties",
+ "required",
+ "minProperties",
+ "maxProperties",
+ "items",
+ "minimum",
+ "maximum",
+ "minLength",
+ "maxLength",
+ "pattern",
+ "default",
+ "anyOf",
+ "propertyOrdering",
+ }
+ )
+
+ def _resolve_local_schema_ref(
+ root: Optional[dict[str, Any]], ref: str
+ ) -> Optional[Any]:
+ # Walk a `#/foo/bar` JSON pointer against the schema root.
+ # Returns None if the pointer doesn't resolve to a dict, so
+ # the caller can fall back to the unresolved node.
+ if not isinstance(root, dict) or not isinstance(ref, str):
+ return None
+ if not ref.startswith("#/"):
+ return None
+ node: Any = root
+ for raw_part in ref[2:].split("/"):
+ if not raw_part:
+ continue
+ part = raw_part.replace("~1", "/").replace("~0", "~")
+ if not isinstance(node, dict) or part not in node:
+ return None
+ node = node[part]
+ return node
+
+ def _sanitize_gemini_schema(
+ node: Any,
+ root: Optional[dict[str, Any]] = None,
+ _seen_refs: Optional[frozenset[str]] = None,
+ ) -> Any:
+ # Recursively filter to Gemini's OpenAPI 3.0 subset. At a
+ # Schema-keyword dict layer we drop keys not in the
+ # allowlist; under `properties` the keys are user-defined
+ # field names and the values are themselves Schemas; under
+ # `items` / `anyOf` the values are also Schemas.
+ # OpenAI strict tools commonly use JSON Schema's
+ # `"type": ["string", "null"]` form for nullable fields;
+ # Gemini's OpenAPI Schema uses `"type": "string"` plus
+ # `"nullable": true`. Translate that here.
+ if root is None and isinstance(node, dict):
+ root = node
+ if _seen_refs is None:
+ _seen_refs = frozenset()
+ if isinstance(node, dict):
+ # Pydantic / OpenAI strict tools commonly hoist nested
+ # object schemas into `$defs` and reference them via
+ # `{"$ref": "#/$defs/Address"}`. Gemini's OpenAPI subset
+ # has no $ref and drops anything not in the allowlist,
+ # so the referenced shape would vanish if we didn't
+ # inline it here. Recurse into the resolved target with
+ # local siblings overriding the reference (normal JSON
+ # Schema composition), guarding against ref cycles.
+ _ref = node.get("$ref")
+ if isinstance(_ref, str):
+ if _ref in _seen_refs:
+ return {}
+ _target = _resolve_local_schema_ref(root, _ref)
+ if isinstance(_target, dict):
+ _merged = {
+ **_target,
+ **{k: v for k, v in node.items() if k != "$ref"},
+ }
+ return _sanitize_gemini_schema(
+ _merged, root, _seen_refs | {_ref}
+ )
+ cleaned: dict[str, Any] = {}
+ _nullable_from_union = False
+ _flattened_type: Optional[str] = None
+ _union_any_of: Optional[list[dict[str, Any]]] = None
+ _raw_type = node.get("type")
+ if isinstance(_raw_type, list):
+ _non_null = [t for t in _raw_type if t != "null"]
+ if len(_non_null) < len(_raw_type):
+ _nullable_from_union = True
+ if len(_non_null) == 1:
+ _flattened_type = _non_null[0]
+ elif len(_non_null) > 1:
+ # Preserve multi-type unions as anyOf; flattening
+ # to the first non-null type silently drops the
+ # other branches and changes the tool contract.
+ _union_any_of = [
+ {"type": _t} for _t in _non_null if isinstance(_t, str)
+ ]
+ for _k, _v in node.items():
+ if _k == "type" and isinstance(_v, list):
+ # Handled below via _flattened_type.
+ continue
+ if _k not in _GEMINI_ALLOWED_SCHEMA_KEYS:
+ continue
+ if _k == "properties" and isinstance(_v, dict):
+ cleaned[_k] = {
+ _name: _sanitize_gemini_schema(_subschema, root, _seen_refs)
+ for _name, _subschema in _v.items()
+ }
+ elif _k == "items":
+ cleaned[_k] = _sanitize_gemini_schema(_v, root, _seen_refs)
+ elif _k == "anyOf" and isinstance(_v, list):
+ # Optional[X] / Union[A, B, None]: Pydantic emits
+ # `anyOf: [..., {"type":"null"}]`. Gemini's
+ # OpenAPI subset rejects `"type": "null"` inside
+ # anyOf, so drop the null variant and surface it
+ # via `nullable: true`. If exactly one non-null
+ # branch remains, collapse it inline; otherwise
+ # keep the slim anyOf and mark the field
+ # nullable.
+ _saw_null = any(
+ isinstance(_entry, dict) and _entry.get("type") == "null"
+ for _entry in _v
+ )
+ _non_null_entries = [
+ _entry
+ for _entry in _v
+ if not (
+ isinstance(_entry, dict)
+ and _entry.get("type") == "null"
+ )
+ ]
+ if len(_non_null_entries) == 1 and _saw_null:
+ _inner = _sanitize_gemini_schema(
+ _non_null_entries[0], root, _seen_refs
+ )
+ if isinstance(_inner, dict):
+ for _ik, _iv in _inner.items():
+ cleaned.setdefault(_ik, _iv)
+ cleaned.setdefault("nullable", True)
+ else:
+ cleaned[_k] = [
+ _sanitize_gemini_schema(_entry, root, _seen_refs)
+ for _entry in _non_null_entries
+ ]
+ if _saw_null:
+ cleaned.setdefault("nullable", True)
+ elif _k in ("required", "enum", "propertyOrdering"):
+ # Lists of plain strings; copy verbatim.
+ cleaned[_k] = _v
+ else:
+ cleaned[_k] = _v
+ if _union_any_of is not None and "anyOf" not in cleaned:
+ cleaned["anyOf"] = [
+ _sanitize_gemini_schema(_s, root, _seen_refs)
+ for _s in _union_any_of
+ ]
+ elif _flattened_type is not None:
+ cleaned["type"] = _flattened_type
+ if _nullable_from_union and "nullable" not in cleaned:
+ cleaned["nullable"] = True
+ return cleaned
+ return node
+
+ function_declarations: list[dict[str, Any]] = []
+ if tools and text_tools_allowed and not _tool_choice_disabled:
+ for _tool in tools:
+ if not isinstance(_tool, dict) or _tool.get("type") != "function":
+ continue
+ _fn = _tool.get("function")
+ if not isinstance(_fn, dict) or not _fn.get("name"):
+ continue
+ _decl: dict[str, Any] = {
+ "name": _fn["name"],
+ "description": _fn.get("description") or "",
+ }
+ _params = _fn.get("parameters")
+ if isinstance(_params, dict):
+ _decl["parameters"] = _sanitize_gemini_schema(_params)
+ function_declarations.append(_decl)
+ if function_declarations:
+ tools_array.append({"functionDeclarations": function_declarations})
+ if tools_array:
+ body["tools"] = tools_array
+ # Tool-choice mapping: OpenAI "auto"/"none"/"required"/{name=...}
+ # -> Gemini toolConfig.functionCallingConfig.mode + allowedFunctionNames.
+ if tool_choice is not None and function_declarations and text_tools_allowed:
+ _mode: Optional[str] = None
+ _allowed: Optional[list[str]] = None
+ if isinstance(tool_choice, str):
+ _tc_lc = tool_choice.strip().lower()
+ if _tc_lc == "auto":
+ _mode = "AUTO"
+ elif _tc_lc == "none":
+ _mode = "NONE"
+ elif _tc_lc in ("required", "any"):
+ _mode = "ANY"
+ elif (
+ isinstance(tool_choice, dict) and tool_choice.get("type") == "function"
+ ):
+ _fn_pick = tool_choice.get("function") or {}
+ _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None
+ if isinstance(_name, str) and _name:
+ _mode = "ANY"
+ _allowed = [_name]
+ if _mode is not None:
+ _fcc: dict[str, Any] = {"mode": _mode}
+ if _allowed:
+ _fcc["allowedFunctionNames"] = _allowed
+ body["toolConfig"] = {"functionCallingConfig": _fcc}
+
+ # Prompt caching. The Gemini caching contract is "create a
+ # CachedContent resource, then pass its name on
+ # `cachedContent`". The cache itself is created out of band by
+ # the caller via POST /cachedContents; here we forward an
+ # explicit cache id when the dispatcher hands us one (a string
+ # value on enable_prompt_caching means "use this cache name").
+ # https://ai.google.dev/gemini-api/docs/caching
+ if isinstance(enable_prompt_caching, str) and enable_prompt_caching:
+ body["cachedContent"] = enable_prompt_caching
+
+ # Model id is already validated at the top of _stream_gemini so
+ # we never reach a path-traversed URL segment here.
+ url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse"
+ completion_id = f"chatcmpl-gemini-{model.replace('/', '-')}"
+
+ logger.info(
+ "Proxying Gemini streamGenerateContent to %s (model=%s, "
+ "tools=%s, image=%s)",
+ url,
+ model,
+ [list(t.keys())[0] for t in tools_array] if tools_array else [],
+ is_image_model,
+ )
+
+ def _emit_tool_event(payload: dict[str, Any]) -> str:
+ _stamp_server_tool_marker(payload)
+ chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {},
+ "finish_reason": None,
+ }
+ ],
+ "_toolEvent": payload,
+ }
+ return f"data: {_json.dumps(chunk)}"
+
+ def _text_chunk(
+ text: str, extra_content: Optional[dict[str, Any]] = None
+ ) -> str:
+ delta: dict[str, Any] = {"content": text}
+ if extra_content:
+ delta["extra_content"] = extra_content
+ chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": delta,
+ "finish_reason": None,
+ }
+ ],
+ }
+ return f"data: {_json.dumps(chunk)}"
+
+ def _gemini_part_extra(part: dict[str, Any]) -> Optional[dict[str, Any]]:
+ """Return ``{"google": {"thought_signature": ...}}`` when the
+ Gemini stream part carries a `thoughtSignature` we need to
+ replay on a follow-up turn (Gemini 3 image editing + tool
+ contexts both require an exact signature echo)."""
+ sig = part.get("thoughtSignature") or part.get("thought_signature")
+ if isinstance(sig, str) and sig:
+ return {"google": {"thought_signature": sig}}
+ return None
+
+ # Gemini finish reasons -> OpenAI vocabulary. Reference:
+ # https://ai.google.dev/api/rest/v1beta/Candidate#FinishReason
+ _finish_reason_map: dict[str, Optional[str]] = {
+ "STOP": "stop",
+ "MAX_TOKENS": "length",
+ "SAFETY": "content_filter",
+ "RECITATION": "content_filter",
+ "PROHIBITED_CONTENT": "content_filter",
+ "BLOCKLIST": "content_filter",
+ "MALFORMED_FUNCTION_CALL": "stop",
+ "OTHER": "stop",
+ "FINISH_REASON_UNSPECIFIED": None,
+ }
+
+ last_usage: Optional[dict[str, Any]] = None
+ emitted_function_call_ids: set[str] = set()
+ # True once any Gemini functionCall part has been emitted so the
+ # final finish_reason swaps STOP -> tool_calls (matches the
+ # OpenAI Chat Completions contract; an OAI client that sees a
+ # tool_calls delta followed by finish_reason="stop" never
+ # executes the tool).
+ emitted_any_function_call = False
+ # web_search_active drives the tool_start / tool_end envelope.
+ # Track on whether `googleSearch` was actually forwarded above,
+ # not the raw caller intent -- image-mode requests filter the
+ # tool out, and emitting a phantom "search complete" card on a
+ # turn where Gemini was never told to search confuses the UI.
+ web_search_active = any("googleSearch" in t for t in tools_array)
+ web_search_tool_id = "gemini_web_search"
+ web_search_tool_started = False
+ web_search_tool_ended = False
+ web_search_citations: list[dict[str, str]] = []
+ # Tracks the tool_call_id minted on the most recent
+ # executableCode part so the matching codeExecutionResult can
+ # close out the same envelope. None between rounds.
+ gemini_code_exec_pending_id: Optional[str] = None
+ # The most recently emitted code_execution id + result text. Kept
+ # *after* the tool_end so a following inline image (matplotlib
+ # plot rendered by codeExecution) can attach to the same card
+ # via a `__IMAGES__:` marker instead of spawning a separate
+ # image_generation event.
+ last_code_exec_tool_id: Optional[str] = None
+ last_code_exec_result_text: str = ""
+
+ try:
+ async with _http_client.stream(
+ "POST",
+ url,
+ json = body,
+ headers = self._auth_headers(),
+ timeout = self._stream_timeout,
+ ) as response:
+ if response.status_code != 200:
+ error_body = await response.aread()
+ error_text = error_body.decode("utf-8", errors = "replace")
+ logger.error(
+ "Gemini returned %d: %s",
+ response.status_code,
+ error_text[:500],
+ )
+ yield _error_sse_line(
+ response.status_code, error_text, self.provider_type
+ )
+ return
+
+ if web_search_active:
+ yield _emit_tool_event(
+ {
+ "type": "tool_start",
+ "tool_name": "web_search",
+ "tool_call_id": web_search_tool_id,
+ "arguments": {},
+ }
+ )
+ web_search_tool_started = True
+
+ # NOTE: same manual __anext__ loop pattern as the other
+ # streaming helpers (see stream_chat_completion for the
+ # Python 3.13 + httpcore 1.0.x GeneratorExit ordering).
+ lines_gen = response.aiter_lines().__aiter__()
+ final_finish_reason: Optional[str] = None
+ try:
+ while True:
+ try:
+ line = await lines_gen.__anext__()
+ except StopAsyncIteration:
+ break
+ if not line.strip():
+ continue
+ if not line.startswith("data:"):
+ continue
+ data_str = line[len("data:") :].strip()
+ if not data_str or data_str == "[DONE]":
+ continue
+ try:
+ event = _json.loads(data_str)
+ except Exception:
+ logger.warning(
+ "Gemini: failed to parse SSE chunk: %s",
+ data_str[:200],
+ )
+ continue
+ if not isinstance(event, dict):
+ continue
+
+ # Latch usageMetadata across deltas -- the final
+ # fragment carries the complete totals.
+ usage_meta = event.get("usageMetadata")
+ if isinstance(usage_meta, dict):
+ last_usage = usage_meta
+
+ # Prompt-level safety block: Gemini ships zero
+ # candidates plus a `promptFeedback.blockReason`
+ # (e.g. SAFETY). The downstream OAI client would
+ # otherwise see an empty successful assistant
+ # response. Surface as a content_filter error
+ # event so the UI can render the block reason.
+ prompt_feedback = event.get("promptFeedback")
+ if isinstance(prompt_feedback, dict) and prompt_feedback.get(
+ "blockReason"
+ ):
+ block_reason = str(prompt_feedback.get("blockReason"))
+ # Close out the synthetic web_search start so
+ # the UI does not show a spinner stuck on
+ # "searching..." after the error toast lands.
+ if (
+ web_search_active
+ and web_search_tool_started
+ and not web_search_tool_ended
+ ):
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": web_search_tool_id,
+ "result": (
+ "(search aborted: Gemini blocked "
+ f"prompt: {block_reason})"
+ ),
+ }
+ )
+ web_search_tool_ended = True
+ yield _error_sse_line(
+ 400,
+ f"Gemini blocked prompt: {block_reason}",
+ self.provider_type,
+ )
+ return
+
+ candidates = event.get("candidates") or []
+ if not isinstance(candidates, list):
+ continue
+ for cand in candidates:
+ if not isinstance(cand, dict):
+ continue
+ # Citations / grounding metadata.
+ # `groundingMetadata.groundingChunks[].web`
+ # carries `uri` + `title`. Collect for the
+ # tool_end emission at stream close.
+ gm = cand.get("groundingMetadata")
+ if isinstance(gm, dict) and web_search_active:
+ chunks_list = gm.get("groundingChunks") or []
+ if isinstance(chunks_list, list):
+ for ch in chunks_list:
+ if not isinstance(ch, dict):
+ continue
+ web = ch.get("web") or {}
+ if not isinstance(web, dict):
+ continue
+ u = web.get("uri") or ""
+ if not u or not isinstance(u, str):
+ continue
+ if any(
+ c["url"] == u for c in web_search_citations
+ ):
+ continue
+ web_search_citations.append(
+ {
+ "url": u,
+ "title": (web.get("title") or u),
+ "snippet": "",
+ }
+ )
+
+ content_obj = cand.get("content") or {}
+ parts = (
+ content_obj.get("parts")
+ if isinstance(content_obj, dict)
+ else None
+ )
+ if isinstance(parts, list):
+ for part in parts:
+ if not isinstance(part, dict):
+ continue
+ # Text delta. Stow part-level
+ # `thoughtSignature` on the delta so
+ # Gemini 3 turns that need an exact
+ # signature echo round-trip cleanly.
+ text = part.get("text")
+ _part_extra = _gemini_part_extra(part)
+ if isinstance(text, str) and text:
+ yield _text_chunk(
+ text,
+ extra_content = _part_extra,
+ )
+ elif _part_extra is not None and not any(
+ k in part
+ for k in (
+ "functionCall",
+ "executableCode",
+ "codeExecutionResult",
+ "inlineData",
+ )
+ ):
+ # Empty-content part carrying a
+ # thoughtSignature: emit an empty delta
+ # so the signature is preserved.
+ yield _text_chunk(
+ "",
+ extra_content = _part_extra,
+ )
+ # functionCall -> OpenAI tool_calls
+ # delta envelope.
+ fc = part.get("functionCall")
+ if isinstance(fc, dict):
+ fc_name = fc.get("name") or ""
+ fc_args = fc.get("args") or {}
+ fc_id = (
+ fc.get("id")
+ or f"call_{fc_name}_{time.time_ns()}"
+ )
+ if fc_id in emitted_function_call_ids:
+ continue
+ emitted_function_call_ids.add(fc_id)
+ # Each distinct functionCall in an
+ # assistant turn needs its own
+ # tool_calls[*].index. Consumers
+ # that reassemble tool_calls by
+ # index collapse all calls onto
+ # the same slot when this is
+ # hardcoded to 0, breaking
+ # parallel/multi-tool turns.
+ tc_index = len(emitted_function_call_ids) - 1
+ tool_call_delta: dict[str, Any] = {
+ "index": tc_index,
+ "id": fc_id,
+ "type": "function",
+ "function": {
+ "name": fc_name,
+ "arguments": _json.dumps(fc_args),
+ },
+ }
+ # Gemini 3 function-calling: the
+ # part-level `thoughtSignature`
+ # must be echoed back on the
+ # next turn or the model rejects
+ # the tool-result envelope. Stow
+ # it on `extra_content.google`
+ # so the frontend can persist it
+ # and our outbound translator
+ # (below) can replay it.
+ thought_sig = part.get(
+ "thoughtSignature"
+ ) or part.get("thought_signature")
+ if isinstance(thought_sig, str) and thought_sig:
+ tool_call_delta["extra_content"] = {
+ "google": {
+ "thought_signature": thought_sig,
+ }
+ }
+ emitted_any_function_call = True
+ tool_chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {
+ "tool_calls": [tool_call_delta]
+ },
+ "finish_reason": None,
+ }
+ ],
+ }
+ yield f"data: {_json.dumps(tool_chunk)}"
+ # executableCode + codeExecutionResult
+ # parts surface as the standard
+ # code_execution tool_start/tool_end
+ # envelope (same shape OpenAI and
+ # Anthropic emit) so the chat
+ # adapter can render Gemini sandbox
+ # output through CodeExecutionToolUI.
+ # https://ai.google.dev/gemini-api/docs/code-execution
+ exec_code = part.get("executableCode")
+ if isinstance(exec_code, dict):
+ code_str = exec_code.get("code") or ""
+ if code_str:
+ code_tool_id = (
+ exec_code.get("id")
+ or f"gemini_code_exec_{time.time_ns()}"
+ )
+ gemini_code_exec_pending_id = code_tool_id
+ # Stow the raw Gemini part so
+ # follow-up turns can replay
+ # the native `executableCode`
+ # (Gemini rejects a generic
+ # functionCall echo for code
+ # execution history).
+ _exec_thought_sig = part.get(
+ "thoughtSignature"
+ ) or part.get("thought_signature")
+ # Per-part thoughtSignature stays
+ # bound to its own part (Gemini 3
+ # rejects shared signatures).
+ _exec_part_entry: dict[str, Any] = {
+ "executableCode": exec_code,
+ }
+ if (
+ isinstance(_exec_thought_sig, str)
+ and _exec_thought_sig
+ ):
+ _exec_part_entry["thoughtSignature"] = (
+ _exec_thought_sig
+ )
+ _exec_native: dict[str, Any] = {
+ "parts": [_exec_part_entry],
+ }
+ yield _emit_tool_event(
+ {
+ "type": "tool_start",
+ "tool_name": "code_execution",
+ "tool_call_id": code_tool_id,
+ "arguments": {
+ "kind": "code_execution",
+ "language": (
+ (
+ exec_code.get(
+ "language"
+ )
+ or "PYTHON"
+ ).lower()
+ ),
+ "code": code_str,
+ "google": {
+ "native_part": _exec_native,
+ },
+ },
+ }
+ )
+ exec_result = part.get("codeExecutionResult")
+ if isinstance(exec_result, dict):
+ outcome = exec_result.get("outcome") or ""
+ output = exec_result.get("output") or ""
+ # Gemini returns
+ # OUTCOME_OK / OUTCOME_FAILED /
+ # OUTCOME_DEADLINE_EXCEEDED. Treat
+ # non-OK outcomes as stderr so the
+ # UI surfaces the error.
+ if outcome and outcome != "OUTCOME_OK":
+ result_text = (
+ f"[{outcome}]\n{output}".rstrip()
+ )
+ else:
+ result_text = output
+ # Pair tool_end with the most recent
+ # executableCode tool_start; fall back
+ # to exec_result.id then a fresh id.
+ pair_id = (
+ gemini_code_exec_pending_id
+ or exec_result.get("id")
+ or f"gemini_code_exec_{time.time_ns()}"
+ )
+ if gemini_code_exec_pending_id is None:
+ yield _emit_tool_event(
+ {
+ "type": "tool_start",
+ "tool_name": "code_execution",
+ "tool_call_id": pair_id,
+ "arguments": {
+ "kind": "code_execution",
+ "code": "",
+ },
+ }
+ )
+ _result_thought_sig = part.get(
+ "thoughtSignature"
+ ) or part.get("thought_signature")
+ _result_part_entry: dict[str, Any] = {
+ "codeExecutionResult": exec_result,
+ }
+ if (
+ isinstance(_result_thought_sig, str)
+ and _result_thought_sig
+ ):
+ _result_part_entry["thoughtSignature"] = (
+ _result_thought_sig
+ )
+ _result_native: dict[str, Any] = {
+ "parts": [_result_part_entry],
+ }
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": pair_id,
+ "result": result_text,
+ "google": {
+ "native_part": _result_native,
+ },
+ }
+ )
+ last_code_exec_tool_id = pair_id
+ last_code_exec_result_text = result_text
+ gemini_code_exec_pending_id = None
+ # inlineData: either a Nano Banana
+ # generation (own card) or a sandbox
+ # plot attached to the code_execution
+ # card via the __IMAGES__: marker.
+ inline = part.get("inlineData")
+ if isinstance(inline, dict):
+ b64 = inline.get("data") or ""
+ mime = inline.get("mimeType") or "image/png"
+ if b64:
+ image_uri = f"data:{mime};base64,{b64}"
+ attached_to_code_exec = (
+ not is_image_model
+ and last_code_exec_tool_id is not None
+ and bool(enabled_tools)
+ and "code_execution"
+ in (enabled_tools or [])
+ )
+ if attached_to_code_exec:
+ updated_result = (
+ last_code_exec_result_text
+ + "\n__IMAGES__:"
+ + _json.dumps([image_uri])
+ )
+ # Stow inlineData so a follow-up
+ # turn can replay the plot with
+ # its per-part thoughtSignature.
+ _plot_thought_sig = part.get(
+ "thoughtSignature"
+ ) or part.get("thought_signature")
+ _plot_part_entry: dict[str, Any] = {
+ "inlineData": {
+ "mimeType": mime,
+ "data": b64,
+ },
+ }
+ if (
+ isinstance(_plot_thought_sig, str)
+ and _plot_thought_sig
+ ):
+ _plot_part_entry[
+ "thoughtSignature"
+ ] = _plot_thought_sig
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": (
+ last_code_exec_tool_id
+ ),
+ "result": updated_result,
+ "google": {
+ "native_part": {
+ "parts": [
+ _plot_part_entry
+ ],
+ },
+ },
+ }
+ )
+ last_code_exec_result_text = (
+ updated_result
+ )
+ else:
+ img_id = f"img_{time.time_ns()}"
+ yield _emit_tool_event(
+ {
+ "type": "tool_start",
+ "tool_name": "image_generation",
+ "tool_call_id": img_id,
+ "arguments": {
+ "kind": "image",
+ "prompt": "",
+ },
+ }
+ )
+ # Gemini 3 image edit needs
+ # the prior thoughtSignature
+ # echoed on the inline image part.
+ _img_thought_sig = part.get(
+ "thoughtSignature"
+ ) or part.get("thought_signature")
+ _img_tool_end: dict[str, Any] = {
+ "type": "tool_end",
+ "tool_call_id": img_id,
+ "result": "",
+ "image_b64": b64,
+ "image_mime": mime,
+ }
+ # Stow inlineData so multi-turn
+ # edits replay the original
+ # image as native history.
+ _img_part_entry: dict[str, Any] = {
+ "inlineData": {
+ "mimeType": mime,
+ "data": b64,
+ },
+ }
+ if (
+ isinstance(_img_thought_sig, str)
+ and _img_thought_sig
+ ):
+ _img_part_entry[
+ "thoughtSignature"
+ ] = _img_thought_sig
+ _img_native: dict[str, Any] = {
+ "parts": [_img_part_entry],
+ }
+ _img_google: dict[str, Any] = {
+ "native_part": _img_native,
+ }
+ if (
+ isinstance(_img_thought_sig, str)
+ and _img_thought_sig
+ ):
+ _img_google["thought_signature"] = (
+ _img_thought_sig
+ )
+ _img_tool_end["google"] = _img_google
+ yield _emit_tool_event(_img_tool_end)
+ finish_reason = cand.get("finishReason")
+ if isinstance(finish_reason, str):
+ mapped = _finish_reason_map.get(finish_reason, "stop")
+ if mapped is not None:
+ final_finish_reason = mapped
+
+ # End-of-stream emission order: web_search tool_end
+ # (with citations) -> finish_reason chunk -> usage
+ # chunk -> [DONE]. Matches the Anthropic / OpenAI
+ # helpers' contract so the frontend handler does
+ # not need provider-specific ordering knowledge.
+ if (
+ web_search_active
+ and web_search_tool_started
+ and not web_search_tool_ended
+ ):
+ blocks: list[str] = []
+ for cit in web_search_citations:
+ line_out = f"Title: {cit['title']}\nURL: {cit['url']}"
+ if cit.get("snippet"):
+ line_out += f"\nSnippet: {cit['snippet']}"
+ blocks.append(line_out)
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": web_search_tool_id,
+ "result": (
+ "\n---\n".join(blocks)
+ if blocks
+ else "(search complete)"
+ ),
+ }
+ )
+ web_search_tool_ended = True
+
+ if final_finish_reason:
+ # OpenAI clients trigger tool execution when
+ # finish_reason="tool_calls". Gemini emits
+ # "STOP" even when the turn was a pure
+ # functionCall request, so override after the
+ # fact to match the OAI contract.
+ if emitted_any_function_call and final_finish_reason == "stop":
+ final_finish_reason = "tool_calls"
+ finish_chunk = {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {},
+ "finish_reason": final_finish_reason,
+ }
+ ],
+ }
+ yield f"data: {_json.dumps(finish_chunk)}"
+
+ # Map Gemini usageMetadata onto OpenAI include_usage.
+ # thoughtsTokenCount is billed output too — fold it in
+ # so cost calculators don't undercount.
+ if isinstance(last_usage, dict):
+ thought_tokens = last_usage.get("thoughtsTokenCount") or 0
+ candidate_tokens = last_usage.get("candidatesTokenCount") or 0
+ prompt_tokens = last_usage.get("promptTokenCount") or 0
+ # Gemini bills tool-call prompt slices separately
+ # via `toolUsePromptTokenCount`. Fold into input
+ # so total_tokens does not undercount tool turns.
+ tool_use_prompt_tokens = (
+ last_usage.get("toolUsePromptTokenCount") or 0
+ )
+ translated_usage = {
+ "input_tokens": prompt_tokens + tool_use_prompt_tokens,
+ "output_tokens": candidate_tokens + thought_tokens,
+ "input_tokens_details": {
+ "cached_tokens": (
+ last_usage.get("cachedContentTokenCount") or 0
+ ),
+ "tool_use_prompt_tokens": tool_use_prompt_tokens,
+ },
+ "output_tokens_details": {
+ "reasoning_tokens": thought_tokens,
+ },
+ }
+ usage_line = _build_usage_chunk(
+ completion_id, "openai", translated_usage
+ )
+ if usage_line:
+ yield usage_line
+
+ yield "data: [DONE]"
+ finally:
+ # Close response first so lines_gen.aclose() becomes
+ # a no-op (avoids the httpcore 1.0 GeneratorExit
+ # path and the aclose-never-awaited RuntimeWarning).
+ await response.aclose()
+ await lines_gen.aclose()
+
+ except httpx.ConnectError as exc:
+ logger.error("Connection error to %s: %s", self.provider_type, exc)
+ if web_search_tool_started and not web_search_tool_ended:
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": web_search_tool_id,
+ "result": f"(search aborted: connection error: {exc})",
+ }
+ )
+ web_search_tool_ended = True
+ yield _error_sse_line(
+ 502,
+ f"Failed to connect to {self.provider_type}: {exc}",
+ self.provider_type,
+ )
+ except httpx.ReadTimeout as exc:
+ logger.error("Read timeout from %s: %s", self.provider_type, exc)
+ if web_search_tool_started and not web_search_tool_ended:
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": web_search_tool_id,
+ "result": "(search aborted: read timeout)",
+ }
+ )
+ web_search_tool_ended = True
+ yield _error_sse_line(
+ 504,
+ f"Timeout waiting for {self.provider_type} response",
+ self.provider_type,
+ )
+ except httpx.HTTPError as exc:
+ logger.error("HTTP error from %s: %s", self.provider_type, exc)
+ if web_search_tool_started and not web_search_tool_ended:
+ yield _emit_tool_event(
+ {
+ "type": "tool_end",
+ "tool_call_id": web_search_tool_id,
+ "result": f"(search aborted: transport error: {exc})",
+ }
+ )
+ web_search_tool_ended = True
+ yield _error_sse_line(
+ 502,
+ f"Error communicating with {self.provider_type}: {exc}",
+ self.provider_type,
+ )
+
async def _stream_openai_responses(
self,
messages: list[dict[str, Any]],
@@ -2893,6 +4949,8 @@ class ExternalProviderClient:
enable_prompt_caching: Optional[bool] = None,
openai_code_exec_container_id: Optional[str] = None,
compaction_threshold: Optional[int] = None,
+ tools: Optional[list[dict[str, Any]]] = None,
+ tool_choice: Optional[Any] = None,
) -> AsyncGenerator[str, None]:
"""
Call OpenAI's /v1/responses endpoint and translate its SSE stream back
@@ -2916,6 +4974,12 @@ class ExternalProviderClient:
# translate user/assistant messages into the Responses input shape.
instructions_parts: list[str] = []
input_items: list[dict[str, Any]] = []
+ # When we drop a server-side builtin `function_call` here, the
+ # matching `role="tool"` follow-up must also be dropped --
+ # otherwise the outbound body contains an orphan
+ # `function_call_output` with no matching `function_call`, which
+ # OpenAI Responses can reject or mis-associate.
+ skipped_server_builtin_call_ids: set[str] = set()
openai_replay_items: list[dict[str, Any]] = []
previous_response_id: Optional[str] = None
for msg in messages:
@@ -2932,6 +4996,126 @@ class ExternalProviderClient:
instructions_parts.append(part["text"])
continue
+ # OpenAI Responses uses item-shape history for function
+ # calling: assistant turns that invoked user tools must
+ # serialize each call as a `function_call` input item, and
+ # each role="tool" follow-up as a `function_call_output`
+ # item keyed by the matching `call_id`. Without this the
+ # second turn after a function call sends Chat Completions
+ # shape and Responses 400s the request.
+ if role == "tool":
+ _call_id = msg.get("tool_call_id") or ""
+ # If the matching assistant `function_call` was a
+ # server-side builtin we already dropped, drop the
+ # follow-up too to avoid emitting an orphan
+ # `function_call_output`.
+ if _call_id and _call_id in skipped_server_builtin_call_ids:
+ continue
+ if isinstance(content, list):
+ _flat_parts: list[str] = []
+ for part in content:
+ if part.get("type") == "text" and part.get("text"):
+ _flat_parts.append(part["text"])
+ _output_text = "".join(_flat_parts)
+ else:
+ _output_text = content if isinstance(content, str) else ""
+ if _call_id:
+ input_items.append(
+ {
+ "type": "function_call_output",
+ "call_id": _call_id,
+ "output": _output_text,
+ }
+ )
+ continue
+
+ # Assistant turns that returned tool_calls translate each
+ # call as a `function_call` item (carrying name + JSON
+ # arguments + call_id). Skip builtin server-side cards
+ # (canonical builtin name + `args._server_tool` marker)
+ # which never round-trip as user functions. We require both
+ # checks so a user function literally named `_server_tool`
+ # in its argument schema is not dropped.
+ _tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
+ if role == "assistant" and isinstance(_tool_calls, list):
+ # Preserve the prior `response.output` ordering: the
+ # model's text precedes its function_call items, and
+ # the matching role=tool follow-up arrives AFTER the
+ # call. Without this guard, history replay puts
+ # function_call -> assistant text -> function_call_output,
+ # which can put the tool output after an unrelated
+ # assistant message and confuse multi-turn function
+ # calling.
+ if isinstance(content, str) and content:
+ input_items.append({"role": "assistant", "content": content})
+ elif isinstance(content, list):
+ _asst_parts: list[dict[str, Any]] = []
+ for _part in content:
+ if not isinstance(_part, dict):
+ continue
+ _pt = _part.get("type")
+ if _pt == "text" and _part.get("text"):
+ _asst_parts.append(
+ {
+ "type": "input_text",
+ "text": _part.get("text", ""),
+ }
+ )
+ elif _pt == "image_url":
+ _u = _part.get("image_url", {}).get("url", "")
+ if _u:
+ _asst_parts.append(
+ {"type": "input_image", "image_url": _u}
+ )
+ if _asst_parts:
+ input_items.append(
+ {"role": "assistant", "content": _asst_parts}
+ )
+
+ for _tc in _tool_calls:
+ if not isinstance(_tc, dict):
+ continue
+ _fn = _tc.get("function") or {}
+ if not isinstance(_fn, dict) or not _fn.get("name"):
+ continue
+ _args_raw = _fn.get("arguments") or ""
+ if not isinstance(_args_raw, str):
+ try:
+ _args_raw = _json.dumps(_args_raw)
+ except Exception:
+ _args_raw = ""
+ _fn_name_lc = (_fn.get("name") or "").lower()
+ _is_server_builtin = False
+ if _fn_name_lc in _SERVER_SIDE_BUILTIN_TOOL_NAMES:
+ try:
+ _args_obj = _json.loads(_args_raw) if _args_raw else {}
+ except Exception:
+ _args_obj = None
+ if isinstance(_args_obj, dict):
+ if _args_obj.get("_server_tool") is True:
+ _is_server_builtin = True
+ else:
+ _g = _args_obj.get("google")
+ if isinstance(_g, dict) and isinstance(
+ _g.get("native_part"), dict
+ ):
+ _is_server_builtin = True
+ _call_id_out = _tc.get("id") or f"call_{time.time_ns()}"
+ if _is_server_builtin:
+ skipped_server_builtin_call_ids.add(_call_id_out)
+ continue
+ input_items.append(
+ {
+ "type": "function_call",
+ "call_id": _call_id_out,
+ "name": _fn["name"],
+ "arguments": _args_raw,
+ }
+ )
+ # Assistant text already emitted above (in order) so we
+ # don't fall through to the generic content branches.
+ continue
+
if isinstance(content, str):
input_items.append({"role": role, "content": content})
continue
@@ -3121,44 +5305,15 @@ class ExternalProviderClient:
if max_tokens is not None:
body["max_output_tokens"] = max_tokens
- # Prompt caching on /v1/responses is automatic and free, but the
- # default in-memory policy only survives ~5-10 min of inactivity
- # (up to ~1 hr). Opt into the 24-hour retention policy so a chat
- # left idle overnight still hits the cache on the next turn.
- # Pricing is identical to in_memory per OpenAI's docs.
- #
- # Gated on the base URL because ollama / llama.cpp / "custom"
- # presets all collapse to provider_type="openai" in
- # toExternalBackendProviderType, so they also land in this
- # helper. Those servers expose /v1/responses-shaped routes in
- # some configurations but don't implement
- # prompt_cache_retention — sending the field unconditionally
- # would 400 them. Match the public OpenAI host strictly so the
- # field only goes to OpenAI cloud. Studio's openai model picker
- # is registry-scoped to gpt-5.x / o3 / gpt-4.5, all of which
- # accept this parameter (gpt-5.5+ already defaults to "24h" and
- # rejects "in_memory", so it's a safe no-op there).
- # OpenAI-family cloud: api.openai.com OR Azure OpenAI Foundry
- # (*.openai.azure.com). Both expose the same Responses-API
- # extensions used below -- prompt_cache_retention,
- # context_management compaction, container shell tool -- so
- # treat them uniformly. Non-cloud OpenAI-compatible servers
- # (ollama / llama.cpp / vLLM / "custom" preset) hit /v1/responses
- # without these extensions and would 400 on the unknown body
- # fields, so they intentionally fall outside this gate.
+ # Opt into 24h prompt-cache retention (free, vs the default
+ # ~5-10 min). Gated on the OpenAI cloud host because ollama /
+ # llama.cpp / "custom" presets reach this code path too and
+ # would 400 on the unknown field.
if is_openai_cloud and enable_prompt_caching is not False:
body["prompt_cache_retention"] = "24h"
- # OpenAI server-side context compaction — see
- # https://developers.openai.com/api/docs/guides/compaction
- # When `compaction_threshold` is provided on a cloud OpenAI
- # request, attach `context_management: [{type:"compaction",
- # compact_threshold:N}]` so the API runs server-side
- # compaction when the rendered prompt crosses the threshold.
- # No beta header is required; no dated version pin. The field
- # is silently dropped for non-cloud backends because ollama /
- # llama.cpp / "custom" presets land in this helper and would
- # 400 on an unknown body field.
+ # Server-side context compaction (OpenAI cloud only).
+ # https://developers.openai.com/api/docs/guides/compaction
if (
is_openai_cloud
and compaction_threshold is not None
@@ -3171,55 +5326,92 @@ class ExternalProviderClient:
}
]
- # OpenAI server-side tools — see
- # https://developers.openai.com/api/docs/guides/tools
- # https://developers.openai.com/api/docs/guides/tools-shell
- # The frontend's Search/Code buttons map to the unified
- # enabled_tools shorthand; translate that into the Responses-API
- # tool schema. Other built-in tools (file_search,
- # code_interpreter, image_generation, computer_use_preview) can
- # be added with the same pattern when we surface their toggles.
+ # Map enabled_tools onto Responses-API server tools (cloud only;
+ # local OAI-compat backends 400 on these).
+ # https://developers.openai.com/api/docs/guides/tools
code_execution_enabled_openai = bool(
enabled_tools and "code_execution" in enabled_tools and is_openai_cloud
)
- # OpenAI's image_generation tool is a Responses-API server tool.
- # See https://developers.openai.com/api/docs/guides/tools-image-generation
- # The model picks size / quality / background server-side and
- # delegates rendering to a gpt-image-* family model; the result
- # comes back inline as an `image_generation_call` output item
- # with a base64 image. Available on every gpt-5.x family member
- # plus gpt-4.1 / gpt-4o / o3 per the docs; restrict to cloud
- # OpenAI because the local llama.cpp / ollama backends don't
- # implement it and would 400.
- image_generation_enabled_openai = image_generation_requested
+ image_generation_enabled_openai = bool(
+ enabled_tools and "image_generation" in enabled_tools and is_openai_cloud
+ )
def _openai_image_generation_tool() -> dict[str, Any]:
tool: dict[str, Any] = {"type": "image_generation"}
if image_generation_has_reference:
- # OpenAI's Responses image tool defaults to `auto`. For
- # Studio's explicit follow-up edit flow, force edit mode so
- # the provider uses the previous response / call id as image
- # context instead of treating the text as a fresh generation.
+ # Force edit mode so the prior call id is used as context.
tool["action"] = "edit"
return tool
- if enabled_tools:
- tools_array: list[dict[str, Any]] = []
- if "web_search" in enabled_tools:
+ # Translate Chat-Completions function tools into the Responses
+ # function-tool shape (flattened name/description/parameters).
+ responses_user_function_tools: list[dict[str, Any]] = []
+ if tools:
+ for _tool in tools:
+ if not isinstance(_tool, dict) or _tool.get("type") != "function":
+ continue
+ _fn = _tool.get("function")
+ if not isinstance(_fn, dict) or not _fn.get("name"):
+ continue
+ _entry: dict[str, Any] = {
+ "type": "function",
+ "name": _fn["name"],
+ }
+ if _fn.get("description"):
+ _entry["description"] = _fn["description"]
+ if isinstance(_fn.get("parameters"), dict):
+ _entry["parameters"] = _fn["parameters"]
+ responses_user_function_tools.append(_entry)
+
+ # Translate tool_choice into the Responses shape.
+ _responses_tc_string: Optional[str] = None
+ if isinstance(tool_choice, str):
+ _tc_lc = tool_choice.strip().lower()
+ if _tc_lc in ("auto", "none", "required"):
+ _responses_tc_string = _tc_lc
+ responses_tool_choice: Optional[Any] = None
+ _has_responses_tools = bool(enabled_tools or responses_user_function_tools)
+ if _responses_tc_string is not None and _has_responses_tools:
+ responses_tool_choice = _responses_tc_string
+ elif (
+ tool_choice is not None
+ and responses_user_function_tools
+ and isinstance(tool_choice, dict)
+ and tool_choice.get("type") == "function"
+ ):
+ _fn_pick = tool_choice.get("function") or {}
+ _name = _fn_pick.get("name") if isinstance(_fn_pick, dict) else None
+ if isinstance(_name, str) and _name:
+ responses_tool_choice = {"type": "function", "name": _name}
+
+ _responses_tool_choice_none = _responses_tc_string == "none"
+ # A pinned user function suppresses hosted builtins (privacy +
+ # billing), matching the Gemini / Anthropic / OpenRouter gates.
+ _responses_tool_choice_forced_function = (
+ isinstance(tool_choice, dict)
+ and tool_choice.get("type") == "function"
+ and isinstance(tool_choice.get("function"), dict)
+ and bool(tool_choice["function"].get("name"))
+ )
+ _responses_hosted_builtins_allowed = (
+ not _responses_tool_choice_none
+ and not _responses_tool_choice_forced_function
+ )
+
+ if (
+ enabled_tools or responses_user_function_tools
+ ) and not _responses_tool_choice_none:
+ tools_array: list[dict[str, Any]] = list(responses_user_function_tools)
+ if (
+ _responses_hosted_builtins_allowed
+ and enabled_tools
+ and "web_search" in enabled_tools
+ ):
tools_array.append({"type": "web_search"})
- if code_execution_enabled_openai:
- # `container_auto` lets OpenAI auto-create a fresh
- # container per request; we capture the resulting
- # container_id off the SSE stream and the chat-adapter
- # persists it onto the thread record. Subsequent turns
- # in the same thread pass it back as
- # `openai_code_exec_container_id`, which we translate to
- # `container_reference` here so the model sees
- # filesystem state from prior turns. Container expires
- # after ~20 min of inactivity per OpenAI's default
- # policy — a stale id 400s, the chat-adapter clears it
- # via container_invalidated, and the next turn falls
- # back to auto-create.
+ if _responses_hosted_builtins_allowed and code_execution_enabled_openai:
+ # Reuse the thread's container so filesystem state
+ # persists; auto-create when there isn't one yet. Stale
+ # ids 400 and are cleared via container_invalidated.
shell_env: dict[str, Any]
if openai_code_exec_container_id:
shell_env = {
@@ -3229,10 +5421,12 @@ class ExternalProviderClient:
else:
shell_env = {"type": "container_auto"}
tools_array.append({"type": "shell", "environment": shell_env})
- if image_generation_enabled_openai:
+ if _responses_hosted_builtins_allowed and image_generation_enabled_openai:
tools_array.append(_openai_image_generation_tool())
if tools_array:
body["tools"] = tools_array
+ if responses_tool_choice is not None:
+ body["tool_choice"] = responses_tool_choice
url = f"{self.base_url}/responses"
completion_id = f"chatcmpl-openai-{model.replace('/', '-')}"
@@ -3246,11 +5440,19 @@ class ExternalProviderClient:
first attempt.
"""
attempt_body = dict(body)
- if enabled_tools:
- tools_array_attempt: list[dict[str, Any]] = []
- if "web_search" in enabled_tools:
+ if (
+ enabled_tools or responses_user_function_tools
+ ) and not _responses_tool_choice_none:
+ tools_array_attempt: list[dict[str, Any]] = list(
+ responses_user_function_tools
+ )
+ if (
+ _responses_hosted_builtins_allowed
+ and enabled_tools
+ and "web_search" in enabled_tools
+ ):
tools_array_attempt.append({"type": "web_search"})
- if code_execution_enabled_openai:
+ if _responses_hosted_builtins_allowed and code_execution_enabled_openai:
if container_id_for_this_attempt:
env_attempt: dict[str, Any] = {
"type": "container_reference",
@@ -3261,12 +5463,17 @@ class ExternalProviderClient:
tools_array_attempt.append(
{"type": "shell", "environment": env_attempt}
)
- if image_generation_enabled_openai:
+ if (
+ _responses_hosted_builtins_allowed
+ and image_generation_enabled_openai
+ ):
tools_array_attempt.append(_openai_image_generation_tool())
if tools_array_attempt:
attempt_body["tools"] = tools_array_attempt
else:
attempt_body.pop("tools", None)
+ if responses_tool_choice is not None:
+ attempt_body["tool_choice"] = responses_tool_choice
return attempt_body
def _is_openai_container_expired_error(error_text: str) -> bool:
@@ -3328,52 +5535,26 @@ class ExternalProviderClient:
done_emitted = False
reasoning_open = False
reasoning_emitted = False
- # Latched from response.completed / response.incomplete so
- # the final log can surface input_tokens_details.cached_tokens —
- # the field that proves prompt_cache_retention="24h" is
- # actually hitting OpenAI's cache instead of recomputing
- # the prefix every turn.
+ # Per-call function-tool indexing; distinct slots so
+ # parallel calls don't collide on delta.tool_calls[].index.
+ saw_function_call = False
+ function_call_index = 0
+ # Latched from response.completed/incomplete; surfaces
+ # input_tokens_details.cached_tokens to prove cache hits.
last_usage: Optional[dict[str, Any]] = None
- # Per-call state for OpenAI's server-side web_search tool. Mapped
- # back into our local _toolEvent shape so the existing chat-UI
- # renderer surfaces web_search the same way it does for local
- # tool calls: a "Searching…" tool-call card, then a `tool_end`
- # carrying citations formatted as
- # Title: …\nURL: …\nSnippet: …\n---\n…
- # blocks (which the frontend's parseSourcesFromResult lifts
- # into source content parts at end of stream).
- # web_search_calls preserves insertion order so we can apply
- # the aggregated citation list onto the *last* call's
- # tool_end — that's the one the frontend's source-pill
- # extraction reads (parseSourcesFromResult flatMaps every
- # web_search result, so a single non-empty result is enough
- # to surface all sources at message tail).
- # OpenAI emits url_citation annotations on text deltas, not
- # per call — there's no wire field linking a citation back
- # to a specific search invocation. Hence the shared list.
- # web_search_calls: { item_id -> {query} }
+ # web_search state. Citations are emitted on text deltas
+ # (not per call), so the aggregate list is shared and
+ # applied to the LAST web_search tool_end (parseSourcesFromResult
+ # flatmaps every call, one non-empty is enough).
web_search_calls: dict[str, dict[str, Any]] = {}
all_url_citations: list[dict[str, Any]] = []
- # Shell-tool (code execution) state. OpenAI emits
- # `shell_call` items (model requesting a command list)
- # paired with `shell_call_output` items (execution
- # results). We mirror the Anthropic code-execution UX
- # by emitting one `_toolEvent` tool_start per
- # shell_call and one tool_end per shell_call_output;
- # they're linked via `shell_call_output.call_id`
- # matching `shell_call.id`. Items are independent of
- # web_search (different keyed map).
- # shell_calls: { call_id -> {commands, output} }
+ # shell_calls (code execution): { call_id -> {commands, output} }.
+ # shell_call <-> shell_call_output match by call_id; emit
+ # tool_start/tool_end like the Anthropic UX.
shell_calls: dict[str, dict[str, Any]] = {}
- # Container id captured from the response stream. When
- # it differs from the inbound id, emit a synthetic
- # `container_ready` _toolEvent so the frontend can
- # persist it onto the thread record for the next turn.
- # Where OpenAI surfaces it is documented loosely; we
- # probe two known fields (response.container_id on
- # response.completed, item.environment.container_id on
- # shell_call output items) and latch the first one we
- # see.
+ # Container id latched from response.container_id or
+ # item.environment.container_id; emit container_ready
+ # when it differs from the inbound id.
latched_container_id: Optional[str] = None
container_id_emitted = False
current_openai_response_id: Optional[str] = None
@@ -3453,6 +5634,7 @@ class ExternalProviderClient:
return rendered
def _emit_tool_event(payload: dict[str, Any]) -> str:
+ _stamp_server_tool_marker(payload)
chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
@@ -3775,17 +5957,9 @@ class ExternalProviderClient:
f"ws_{len(web_search_calls)}"
)
web_search_calls.setdefault(item_id, {"query": ""})
- # Shell-tool: register the call eagerly so
- # the matching shell_call_output can link
- # back even if `done` arrives out of order.
- # Also probe for container_id on the
- # environment field — when container_auto
- # auto-creates one, this is the first place
- # the new id might surface (OpenAI doesn't
- # promise this in docs, but the field is
- # cheap to scan and lets us emit
- # container_ready earlier than
- # response.completed).
+ # Register shell_call eagerly so out-of-order
+ # output links back. Probe env.container_id
+ # to emit container_ready before response.completed.
if (
isinstance(item, dict)
and item.get("type") == "shell_call"
@@ -3978,24 +6152,9 @@ class ExternalProviderClient:
}
)
elif item.get("type") == "image_generation_call":
- # OpenAI's image_generation tool returns
- # a single output item with the base64
- # PNG/WebP/JPEG on `result` (sometimes
- # `b64_json` depending on output_format).
- # `revised_prompt` is what the gpt-image
- # backbone actually used after refinement
- # of the assistant's request. Emit
- # tool_start + tool_end so the chat card
- # renders the prompt + the generated
- # image inline. The frontend chat-adapter
- # decides how to render the base64 blob
- # (likely an 
)
- # based on the `kind: "image"` hint we
- # set on tool_start arguments.
- # `time_ns()` (nanoseconds) instead of
- # millisecond resolution so synthesised
- # ids stay unique even when two image
- # generations resolve in the same ms.
+ # Base64 image on `result` (or `b64_json`),
+ # `revised_prompt` for the rewritten prompt.
+ # ns-resolution id so concurrent gens stay unique.
raw_item_id = item.get("id")
item_id = raw_item_id or f"img_{time.time_ns()}"
prompt_in = (
@@ -4034,6 +6193,54 @@ class ExternalProviderClient:
"prompt": prompt_in,
}
)
+ elif item.get("type") == "function_call":
+ # Translate to Chat-Completions delta.tool_calls.
+ # https://platform.openai.com/docs/guides/function-calling?api-mode=responses
+ fn_call_id = (
+ item.get("call_id")
+ or item.get("id")
+ or f"call_{time.time_ns()}"
+ )
+ fn_name = item.get("name") or ""
+ fn_args = item.get("arguments") or ""
+ if not isinstance(fn_args, str):
+ try:
+ fn_args = _json.dumps(fn_args)
+ except Exception:
+ fn_args = ""
+ _tc_index = function_call_index
+ function_call_index += 1
+ yield (
+ "data: "
+ + _json.dumps(
+ {
+ "id": completion_id,
+ "object": "chat.completion.chunk",
+ "choices": [
+ {
+ "index": 0,
+ "delta": {
+ "tool_calls": [
+ {
+ "index": _tc_index,
+ "id": fn_call_id,
+ "type": "function",
+ "function": {
+ "name": fn_name,
+ "arguments": (
+ fn_args
+ ),
+ },
+ }
+ ],
+ },
+ "finish_reason": None,
+ }
+ ],
+ }
+ )
+ )
+ saw_function_call = True
elif (
isinstance(event_type, str)
@@ -4167,7 +6374,11 @@ class ExternalProviderClient:
{
"index": 0,
"delta": {},
- "finish_reason": "stop",
+ "finish_reason": (
+ "tool_calls"
+ if saw_function_call
+ else "stop"
+ ),
}
],
}
@@ -4448,11 +6659,70 @@ class ExternalProviderClient:
models = [model for model in raw_models if isinstance(model, dict)]
if not models and self.provider_type == "ollama":
models = await self._list_ollama_native_models()
+ # Gemini's native /v1beta/models returns
+ # {"models": [{"name": "models/gemini-2.5-flash", ...}]}
+ # -- repackage into the OpenAI-compatible shape the rest
+ # of Studio expects so dynamic model discovery works.
+ if not models and self.provider_type == "gemini":
+ models = self._parse_gemini_models(data)
return models
except httpx.HTTPError as exc:
logger.error("Failed to list models from %s: %s", self.provider_type, exc)
raise
+ @staticmethod
+ def _parse_gemini_models(payload: Any) -> list[dict[str, Any]]:
+ """Translate Gemini's native /v1beta/models payload to OpenAI shape.
+
+ Native response:
+ {"models": [{"name": "models/gemini-2.5-flash",
+ "baseModelId": "gemini-2.5-flash",
+ "displayName": "Gemini 2.5 Flash",
+ "supportedGenerationMethods": [...]}]}
+
+ We only keep entries that advertise
+ ``generateContent`` / ``streamGenerateContent`` so the picker
+ does not surface embedding-only models the chat path can't
+ drive.
+ """
+ if not isinstance(payload, dict):
+ return []
+ entries = payload.get("models") or []
+ if not isinstance(entries, list):
+ return []
+ out: list[dict[str, Any]] = []
+ for entry in entries:
+ if not isinstance(entry, dict):
+ continue
+ methods = entry.get("supportedGenerationMethods") or []
+ if (
+ isinstance(methods, list)
+ and methods
+ and not any(
+ m in methods for m in ("generateContent", "streamGenerateContent")
+ )
+ ):
+ continue
+ base_id = entry.get("baseModelId")
+ name = entry.get("name") or ""
+ # ``name`` arrives as ``"models/gemini-2.5-flash"``; the
+ # chat path uses the bare id.
+ short_id = (
+ base_id
+ if isinstance(base_id, str) and base_id
+ else (name.split("/", 1)[1] if "/" in name else name)
+ )
+ if not short_id:
+ continue
+ out.append(
+ {
+ "id": short_id,
+ "owned_by": "google",
+ "display_name": entry.get("displayName") or short_id,
+ }
+ )
+ return out
+
async def _list_ollama_native_models(self) -> list[dict[str, Any]]:
"""Fallback when Ollama's /v1/models returns an empty or null catalog."""
root = self.base_url.removesuffix("/v1").rstrip("/")
@@ -4757,6 +7027,17 @@ def _build_usage_chunk(
"total_tokens": prompt_tokens + completion_tokens,
"prompt_tokens_details": {"cached_tokens": cached},
}
+ # Surface OpenAI Responses / Gemini reasoning-token detail. The
+ # caller pre-populates last_usage["output_tokens_details"] with
+ # at least {"reasoning_tokens": ...}; mirror it into the OAI
+ # `completion_tokens_details` shape so SDKs can render the
+ # hidden-thoughts slice.
+ out_details = last_usage.get("output_tokens_details")
+ if isinstance(out_details, dict) and out_details:
+ usage_block["completion_tokens_details"] = {
+ "reasoning_tokens": out_details.get("reasoning_tokens") or 0,
+ }
+ usage_block["output_tokens_details"] = out_details
chunk = {
"id": completion_id,
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 76234386aa..cce95fc34c 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -17,6 +17,7 @@ import struct
import structlog
from loggers import get_logger
import shutil
+import signal
import socket
import subprocess
import sys
@@ -28,6 +29,12 @@ from urllib.parse import urlparse
import httpx
+from core.inference.llama_server_args import (
+ parse_cache_override,
+ parse_ctx_override,
+ resolve_cache_type_kv,
+ resolve_requested_ctx,
+)
from core.tool_healing import (
_TC_END_TAG_RE,
_TC_FUNC_CLOSE_RE,
@@ -959,9 +966,6 @@ class LlamaCppBackend:
7. llama-server on PATH (system install)
8. ./bin/llama-server (legacy: extracted binary)
"""
- import os
- import sys
-
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
# 1. Env var — direct path to binary
@@ -1232,6 +1236,33 @@ class LlamaCppBackend:
return total
+ @staticmethod
+ def _amd_apu_wants_unified_memory() -> bool:
+ """True only for AMD unified-memory APUs (gfx1150/gfx1151), where
+ GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
+ False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete
+ GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped."""
+ try:
+ import torch
+
+ if getattr(torch.version, "hip", None) is None:
+ return False
+ if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
+ return False
+ for _i in range(torch.cuda.device_count()):
+ try:
+ _arch = (
+ getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "")
+ or ""
+ )
+ except Exception:
+ continue
+ if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
+ return True
+ except Exception:
+ return False
+ return False
+
@staticmethod
def _get_gpu_free_memory() -> list[tuple[int, int]]:
"""Query free memory per GPU.
@@ -1249,8 +1280,6 @@ class LlamaCppBackend:
Returns list of (gpu_index, free_mib) sorted by index. Empty
list if no supported GPU is reachable.
"""
- import os
-
# ── NVIDIA via nvidia-smi ────────────────────────────────────
try:
result = subprocess.run(
@@ -2562,6 +2591,105 @@ class LlamaCppBackend:
# ── Lifecycle ─────────────────────────────────────────────────
+ # GGUF ``general.architecture`` values for diffusion / image models.
+ # llama.cpp proper has no such architectures, so loading one as a chat
+ # model dies with "unknown model architecture: ''". These match
+ # the patched stable-diffusion.cpp / ComfyUI-GGUF enums (LLM_ARCH_FLUX,
+ # LLM_ARCH_QWEN_IMAGE, ...). Unsloth publishes FLUX and Qwen-Image GGUFs
+ # under https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs.
+ # Matched exactly (not as a substring) so a chat arch merely containing a
+ # short token like "wan"/"sd1" (e.g. "taiwan") is not misrouted to Images.
+ _DIFFUSION_ARCHES = frozenset(
+ (
+ "qwen_image",
+ "flux",
+ "sd1",
+ "sdxl",
+ "sd3",
+ "aura",
+ "hidream",
+ "cosmos",
+ "ltxv",
+ "hyvid",
+ "wan",
+ "lumina2",
+ )
+ )
+
+ @staticmethod
+ def _classify_llama_start_failure(
+ output: str,
+ gguf_path: Optional[str],
+ model_identifier: Optional[str],
+ ) -> str:
+ """Explain *why* llama-server failed to start, from its output.
+
+ Several distinct failures all otherwise collapse into the same
+ opaque "invalid GGUF or out of memory" message. The worst case is
+ a diffusion / image GGUF (FLUX, Qwen-Image, ...) loaded as a chat
+ model: the file is perfectly valid and there is plenty of memory,
+ but llama.cpp has no such architecture, so the user is told to free
+ memory that was never the problem (issue #5842). Pick the most
+ specific message the captured output supports.
+ """
+ lowered = (output or "").lower()
+
+ # Detect Ollama source up front so the arch branch can keep the
+ # Ollama hint instead of the generic "unsupported arch" message.
+ gguf = gguf_path or ""
+ is_ollama = (
+ ".studio_links" in gguf
+ or os.sep + "ollama_links" + os.sep in gguf
+ or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf
+ or (model_identifier or "").startswith("ollama/")
+ )
+
+ # "unknown model architecture: ''": diffusion -> Images page,
+ # Ollama -> Ollama hint, else a precise "unsupported" message. Exact
+ # match so chat archs are never misrouted.
+ arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered)
+ if arch_match:
+ arch = arch_match.group(1)
+ if arch in LlamaCppBackend._DIFFUSION_ARCHES:
+ return (
+ f"'{arch}' is a diffusion (image-generation) GGUF, which "
+ "llama-server cannot run as a chat/completion model. Use "
+ "Studio's Images page to generate with local diffusion "
+ "GGUFs such as FLUX and Qwen-Image."
+ )
+ if is_ollama:
+ return (
+ "Some Ollama models do not work with llama.cpp. Try a "
+ "different model, or use this model directly through "
+ "Ollama instead."
+ )
+ return (
+ f"llama.cpp does not support this GGUF's model architecture "
+ f"('{arch}'). The file is valid, but this model type cannot "
+ "be run with llama-server."
+ )
+
+ # Other Ollama compat failures that do not name an arch. Only when
+ # the output shows a GGUF compat issue, not OOM / missing binaries.
+ if is_ollama:
+ gguf_compat_hints = (
+ "key not found",
+ "unknown model architecture",
+ "failed to load model",
+ )
+ if any(h in lowered for h in gguf_compat_hints):
+ return (
+ "Some Ollama models do not work with llama.cpp. Try a "
+ "different model, or use this model directly through "
+ "Ollama instead."
+ )
+
+ # Fallback: genuinely unknown failure (OOM, missing binary, ...).
+ return (
+ "llama-server failed to start. "
+ "Check that the GGUF file is valid and you have enough memory."
+ )
+
def load_model(
self,
*,
@@ -2724,7 +2852,23 @@ class LlamaCppBackend:
# Select GPU(s) based on model size + estimated KV cache.
# Seed safe defaults before GPU probing so the except path
# still has valid state to publish.
- effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0)
+ ctx_override = parse_ctx_override(extra_args)
+ requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
+ cache_override = parse_cache_override(extra_args)
+ cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
+ if ctx_override is not None and ctx_override > 0:
+ logger.info(
+ f"User --ctx-size {ctx_override} honored; "
+ "skipping auto-reduce"
+ )
+ if cache_override is not None:
+ logger.info(
+ f"User --cache-type-k/-v {cache_override} "
+ "honored for KV estimate"
+ )
+ effective_ctx = (
+ requested_ctx if requested_ctx > 0 else (self._context_length or 0)
+ )
max_available_ctx = self._context_length or effective_ctx
gpus: list[tuple[int, int]] = []
try:
@@ -2734,8 +2878,8 @@ class LlamaCppBackend:
# Resolve effective context: 0 means let llama-server use the
# model's native length. Only expand to a known native length
# if metadata is available; otherwise preserve 0 as a sentinel.
- if n_ctx > 0:
- effective_ctx = n_ctx
+ if requested_ctx > 0:
+ effective_ctx = requested_ctx
elif self._context_length is not None:
effective_ctx = self._context_length
else:
@@ -2788,7 +2932,7 @@ class LlamaCppBackend:
# since multi-GPU is slower and the user didn't ask for a
# specific context length.
gpu_indices, use_fit = None, True
- explicit_ctx = n_ctx > 0
+ explicit_ctx = requested_ctx > 0
if gpus and self._can_estimate_kv() and effective_ctx > 0:
# Compute the largest hardware-aware cap from the model's
@@ -2845,7 +2989,7 @@ class LlamaCppBackend:
gpu_indices, use_fit = self._select_gpus(
requested_total, gpus
)
- # No silent shrink: effective_ctx stays == n_ctx.
+ # No silent shrink: effective_ctx stays == requested_ctx.
else:
# Auto context: prefer fewer GPUs, cap context
# to fit. Same headroom threshold as
@@ -2934,7 +3078,7 @@ class LlamaCppBackend:
except Exception as e:
logger.warning(f"GPU selection failed ({e}), using --fit on")
gpu_indices, use_fit = None, True
- effective_ctx = n_ctx # fall back to original
+ effective_ctx = requested_ctx # fall back to original
launch_mmproj_path = self._resolve_launch_mmproj_path(
model_path = model_path,
@@ -3136,6 +3280,14 @@ class LlamaCppBackend:
env = child_env_without_native_path_secret()
binary_dir = str(Path(binary).parent)
+ # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
+ # shared system RAM. setdefault so a user value wins.
+ if self._amd_apu_wants_unified_memory():
+ env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
+ logger.info(
+ "AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1"
+ )
+
if sys.platform == "win32":
# See _build_windows_path_dirs for ordering. #5106.
path_dirs = self._build_windows_path_dirs(
@@ -3145,6 +3297,24 @@ class LlamaCppBackend:
)
existing_path = env.get("PATH", "")
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
+
+ # ROCm: the llama.cpp prebuilt bundles its own rocblas.dll
+ # but NOT the Tensile kernel library files it needs
+ # (rocblas/library/TensileLibrary*.dat + *.hsaco). The
+ # bundled DLL searches relative to its own location by
+ # default (i.e. /rocblas/library/) which does
+ # not exist, causing a silent crash on the first GEMM.
+ # ROCBLAS_TENSILE_LIBPATH overrides that search to point at
+ # the ROCm installation where the kernel files actually are.
+ _hip_path = os.environ.get(
+ "HIP_PATH", os.environ.get("ROCM_PATH", "")
+ )
+ if _hip_path:
+ _rocblas_lib = os.path.join(
+ _hip_path, "bin", "rocblas", "library"
+ )
+ if os.path.isdir(_rocblas_lib):
+ env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
else:
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
# and CUDA runtime libs (libcudart, libcublas, etc.)
@@ -3312,31 +3482,12 @@ class LlamaCppBackend:
# Wait for llama-server to become healthy
if not self._wait_for_health(timeout = 600.0):
self._kill_process()
- _gguf = gguf_path or ""
- _is_ollama = (
- ".studio_links" in _gguf
- or os.sep + "ollama_links" + os.sep in _gguf
- or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf
- or (self._model_identifier or "").startswith("ollama/")
- )
- # Only show the Ollama-specific message when the server
- # output indicates a GGUF compatibility issue, not for
- # unrelated failures like OOM or missing binaries.
- if _is_ollama:
- _output = "\n".join(self._stdout_lines[-50:]).lower()
- _gguf_compat_hints = (
- "key not found",
- "unknown model architecture",
- "failed to load model",
- )
- if any(h in _output for h in _gguf_compat_hints):
- raise RuntimeError(
- "Some Ollama models do not work with llama.cpp. "
- "Try a different model, or use this model directly through Ollama instead."
- )
raise RuntimeError(
- "llama-server failed to start. "
- "Check that the GGUF file is valid and you have enough memory."
+ self._classify_llama_start_failure(
+ "\n".join(self._stdout_lines[-50:]),
+ gguf_path,
+ self._model_identifier,
+ )
)
self._healthy = True
@@ -3853,10 +4004,6 @@ class LlamaCppBackend:
Falls back to pgrep + /proc//exe on Linux when psutil is
not installed.
"""
- import os
- import signal
- import sys
-
try:
# -- Build the ownership allowlist --------------------------------
# Two kinds of matches:
@@ -5076,13 +5223,30 @@ class LlamaCppBackend:
_effective_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
)
- result = execute_tool(
- tool_name,
- arguments,
- cancel_event = cancel_event,
- timeout = _effective_timeout,
- session_id = session_id,
- )
+ # Guard against the model emitting a tool not in the
+ # per-request advertised set: filtered MCP names, a
+ # built-in the caller opted out of, or a stale name
+ # from a prior turn. Mirrors the safetensors loop's
+ # allowed_tool_names check.
+ _allowed = {
+ (t.get("function") or {}).get("name")
+ for t in (tools or [])
+ if (t.get("function") or {}).get("name")
+ }
+ if _allowed and tool_name not in _allowed:
+ result = (
+ f"Error: tool '{tool_name}' is not enabled "
+ "for this request. Use one of the enabled "
+ "tools or provide a final answer."
+ )
+ else:
+ result = execute_tool(
+ tool_name,
+ arguments,
+ cancel_event = cancel_event,
+ timeout = _effective_timeout,
+ session_id = session_id,
+ )
yield {
"type": "tool_end",
diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py
index d8b7eb383e..b299e1ee9e 100644
--- a/studio/backend/core/inference/llama_server_args.py
+++ b/studio/backend/core/inference/llama_server_args.py
@@ -1,46 +1,29 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-"""Validator for user-supplied llama-server pass-through args.
+"""Boundary validator for user-supplied llama-server pass-through args.
-Studio runs llama-server as a managed subprocess and lets callers pass
-extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP:
-``LoadRequest.llama_extra_args``). This module is the boundary that
-rejects only flags Studio fundamentally cannot share with the user --
-model identity, the auth key, and the network endpoint Studio's HTTP
-proxy targets. Anything else passes through.
+Reject only flags Studio manages (model identity, auth, network,
+parallel slots). Everything else (sampling, ``-c``, ``-ngl``,
+``--flash-attn``, ``--cache-type-*``, ``--spec-*``, ``--jinja``, ...)
+is appended after Studio's auto-set flags so llama.cpp's last-wins
+parser lets the user override.
-User-supplied args are appended to ``cmd`` after Studio's auto-set
-flags, so llama.cpp's last-wins CLI parsing makes the user's value
-override the auto-set one. That covers tunable knobs the user might
-reasonably want to override -- ``-c``/``--ctx-size``,
-``-np``/``--parallel``, ``-fa``/``--flash-attn``,
-``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``,
-``--cache-type-k/v``, ``--chat-template-file/-kwargs``,
-``--spec-*``, ``--jinja``/``--no-jinja``,
-``--no-context-shift``/``--context-shift``, sampling params, etc.
-
-Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
+Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
from __future__ import annotations
from typing import Iterable, Optional
-# Each group is the full set of aliases (short + long) for one
-# hard-denied flag, taken from the llama-server README. If llama.cpp
-# adds a new alias for an existing denied flag, extend the relevant
-# group.
-#
-# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl,
-# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*,
-# --chat-template-*, --spec-*) pass through and override Studio's
-# auto-set version via llama.cpp's last-wins CLI parsing.
+# Each group = every alias (short + long) of one hard-denied flag.
+# Extend the matching group when llama.cpp adds a new alias.
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
- # Model identity -- Studio resolves the model from LoadRequest and
- # passes -m / mmproj after downloading from HF if needed. A second
- # -m would point at a different model than the one Studio thinks
- # is loaded.
+ # Parallel slots: owned by typer --parallel; a pass-through would
+ # desync app.state.llama_parallel_slots from llama-server.
+ frozenset({"-np", "--parallel", "--n-parallel"}),
+ # Model identity: Studio resolves it from LoadRequest; a second
+ # -m would load a different model than Studio thinks it loaded.
frozenset({"-m", "--model"}),
frozenset({"-mu", "--model-url"}),
frozenset({"-dr", "--docker-repo"}),
@@ -51,28 +34,21 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"-hft", "--hf-token"}),
frozenset({"-mm", "--mmproj"}),
frozenset({"-mmu", "--mmproj-url"}),
- # Networking -- Studio binds llama-server's port and reverse-proxies
- # HTTP traffic to it. Retargeting host/port/path/prefix would
- # orphan Studio's proxy and the UI would lose the server.
+ # Networking: Studio binds + proxies; retargeting orphans the proxy.
frozenset({"--host"}),
frozenset({"--port"}),
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
- # Auth / TLS -- Studio terminates auth at its own layer; an
- # upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM
- # key, and TLS on llama-server would break the local proxy hop.
+ # Auth / TLS: Studio terminates auth; upstream --api-key / TLS
+ # shadows Studio's key and breaks the proxy hop.
frozenset({"--api-key"}),
frozenset({"--api-key-file"}),
frozenset({"--ssl-key-file"}),
frozenset({"--ssl-cert-file"}),
- # Single-model server -- Studio runs one model per llama-server
- # process and serves its own UI. Enabling multi-model loading or
- # llama-server's built-in web UI changes the surface clients see.
- # ``--webui``/``--no-webui`` are the legacy spelling; current
- # upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions.
- # Keep both so the denylist matches old and new llama-server
- # binaries (Studio's prebuilt vs system-llama.cpp).
+ # Built-in web UI. --webui/--no-webui is the legacy spelling;
+ # upstream renamed to --ui/--no-ui + --ui-*. Keep both so prebuilt
+ # and system llama.cpp binaries both match.
frozenset({"--webui", "--no-webui"}),
frozenset({"--ui", "--no-ui"}),
frozenset({"--ui-config"}),
@@ -82,32 +58,46 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"--models-preset"}),
frozenset({"--models-max"}),
frozenset({"--models-autoload", "--no-models-autoload"}),
+ # Server-mode flips: --embedding / --rerank restrict llama-server to
+ # those endpoints, breaking Studio's /v1/chat/completions hop.
+ frozenset({"--embedding", "--embeddings"}),
+ frozenset({"--rerank", "--reranking"}),
+ # llama-server's own built-in tools flag would silently stack on top
+ # of Studio's --enable-tools / --disable-tools policy resolver.
+ frozenset({"--tools"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
- """Return the flag name for a token, or None if it isn't a flag.
+ """Flag name for ``token``, or None if it isn't a flag.
- Peels ``--key=value`` to the bare ``--key``. Plain numeric values
- like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags;
- llama-server short-form flags always start with a letter.
+ Peels `--key=value` to `--key`, treats `-1` / `-0.5` as values
+ (llama-server shorts always start with a letter), strips
+ whitespace, and normalises attached `-np8` / signed `-np-1` /
+ digit-prefix-junk `-np8x` to `-np`. Mirrors the CLI's
+ `_expand_attached_np_short`.
"""
+ token = token.strip()
if not token.startswith("-") or token in {"-", "--"}:
return None
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
return None
- return token.split("=", 1)[0]
+ name = token.split("=", 1)[0]
+ if len(name) > 3 and name.startswith("-np"):
+ suffix = name[3:]
+ if suffix[0].isdigit() or (
+ len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit()
+ ):
+ return "-np"
+ return name
def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
- """Validate user-supplied llama-server args.
-
- Returns the args as a flat list ready to extend the llama-server
- command. Raises ``ValueError`` (with the offending flag in the
- message) the moment a token resolves to a Studio-managed flag.
- """
+ """Validate user-supplied llama-server args. Returns a flat list
+ ready to extend the llama-server command; raises ``ValueError``
+ naming the offending flag on the first managed token."""
if not args:
return []
out: list[str] = []
@@ -120,19 +110,21 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
f"and cannot be passed as an extra arg"
)
out.append(token)
+ parse_ctx_override(out)
+ parse_cache_override(out)
return out
def is_managed_flag(flag: str) -> bool:
- """True if ``flag`` is a Studio-managed llama-server flag."""
- return flag in _DENYLIST
+ """True if ``flag`` is Studio-managed. Normalises via ``_flag_name``
+ so `-np8` / `--parallel=8` classify like the canonical tokens."""
+ normalised = _flag_name(flag)
+ return normalised is not None and normalised in _DENYLIST
-# Pass-through flags that shadow first-class ``LoadRequest`` fields
-# (max_seq_length, cache_type_kv, speculative_type,
-# chat_template_override). Stripped from inherited extras so they
-# can't last-wins-override an Apply that re-sets the same first-class
-# field.
+# Pass-through flags that shadow first-class LoadRequest fields;
+# stripped from inherited extras so they can't last-wins-override an
+# Apply that re-sets the same field.
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
_CACHE_FLAGS: frozenset[str] = frozenset(
{"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}
@@ -169,14 +161,124 @@ _SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
)
-# Boolean flags inside _SHADOWING_FLAGS that take no value. The
-# value-consuming heuristic in strip_shadowing_flags must skip just the
-# flag for these, never the following token.
+# Shadowing flags that take no value -- strip the flag only, never the
+# following token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja"}
)
+def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
+ """Return the last user-supplied ``-c`` / ``--ctx-size`` value.
+
+ Mirrors llama.cpp's last-wins flag parsing for the one pass-through
+ numeric knob Studio's load-time fit logic needs to see.
+ """
+ if not args:
+ return None
+
+ tokens = [str(a) for a in args]
+ override: Optional[int] = None
+ i, n = 0, len(tokens)
+ while i < n:
+ tok = tokens[i]
+ flag = _flag_name(tok)
+ if flag is None or flag not in _CONTEXT_FLAGS:
+ i += 1
+ continue
+
+ if "=" in tok:
+ raw_value = tok.split("=", 1)[1]
+ i += 1
+ else:
+ if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
+ raise ValueError(
+ f"llama-server flag '{flag}' requires an integer value"
+ )
+ raw_value = tokens[i + 1]
+ i += 2
+
+ try:
+ value = int(str(raw_value).strip())
+ except ValueError as exc:
+ raise ValueError(
+ f"llama-server flag '{flag}' requires an integer value"
+ ) from exc
+ if value < 0:
+ raise ValueError(
+ f"llama-server flag '{flag}' requires a non-negative integer value"
+ )
+ override = value
+
+ return override
+
+
+def resolve_requested_ctx(
+ args: Optional[Iterable[str]],
+ fallback_n_ctx: int,
+) -> int:
+ """Return the context size load_model should treat as requested.
+
+ Single source of truth for the two-line ``ctx_override = parse_ctx_override(...);
+ requested_ctx = ctx_override if ctx_override is not None else n_ctx`` pattern
+ used by ``load_model`` so tests don't have to reimplement the conditional
+ locally and then assert against their own reimplementation.
+ """
+ override = parse_ctx_override(args)
+ return override if override is not None else fallback_n_ctx
+
+
+def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
+ """Return the last-wins cache type if extras pass cache flags.
+
+ Mirrors parse_ctx_override but for cache type. Recognises both -ctk
+ (key) and -ctv (value). When both flags appear, returns the last-wins
+ value, treating key and value cache flags as the same setting because
+ Studio's KV estimate has a single cache_type_kv knob.
+ """
+ if not args:
+ return None
+
+ tokens = [str(a) for a in args]
+ override: Optional[str] = None
+ i, n = 0, len(tokens)
+ while i < n:
+ tok = tokens[i]
+ flag = _flag_name(tok)
+ if flag is None or flag not in _CACHE_FLAGS:
+ i += 1
+ continue
+
+ if "=" in tok:
+ raw_value = tok.split("=", 1)[1]
+ i += 1
+ else:
+ if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
+ raise ValueError(f"llama-server flag '{flag}' requires a value")
+ raw_value = tokens[i + 1]
+ i += 2
+
+ value = str(raw_value).strip()
+ if not value:
+ raise ValueError(f"llama-server flag '{flag}' requires a non-empty value")
+ override = value
+
+ return override
+
+
+def resolve_cache_type_kv(
+ args: Optional[Iterable[str]],
+ fallback_cache_type_kv: Optional[str],
+) -> Optional[str]:
+ """Return the cache type load_model should treat as requested.
+
+ Single source of truth for the cache override conditional used by
+ ``load_model``.
+ """
+ override = parse_cache_override(args)
+ return override if override is not None else fallback_cache_type_kv
+
+
def strip_shadowing_flags(
args: Iterable[str],
*,
@@ -187,14 +289,11 @@ def strip_shadowing_flags(
) -> list[str]:
"""Strip flags that shadow first-class Studio settings.
- Used when the route inherits a previous load's ``llama_extra_args``
- so that an inherited ``-c 4096`` cannot override the current
- request's ``max_seq_length`` (and equivalents for cache /
- speculative / chat template). Each ``strip_*`` flag controls one
- group; the route only strips groups whose corresponding first-class
- field was actually supplied by the caller, so an inherited
- ``--chat-template-file`` survives an Apply that omits both
- ``llama_extra_args`` and ``chat_template_override``.
+ Used when inheriting a previous load's ``llama_extra_args`` so an
+ inherited `-c 4096` can't override the current `max_seq_length`
+ (same for cache / spec / template). Each ``strip_*`` toggle
+ controls one group; the route only strips groups whose first-class
+ field the caller actually supplied.
"""
shadowing: set[str] = set()
if strip_context:
@@ -216,9 +315,8 @@ def strip_shadowing_flags(
out.append(tok)
i += 1
continue
- # Drop this token. Boolean shadowing flags never carry a value;
- # other shadowing flags consume the next token when it isn't a
- # flag and the value isn't already packed as ``--key=value``.
+ # Drop the flag; consume the next token too unless it's
+ # boolean, already inline (`-c=4096`), or another flag.
if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok:
i += 1
elif i + 1 < n and _flag_name(tokens[i + 1]) is None:
diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py
new file mode 100644
index 0000000000..2ed1a630dc
--- /dev/null
+++ b/studio/backend/core/inference/mcp_client.py
@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import shlex
+import sys
+from typing import Any, Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+MCP_TOOL_PREFIX = "mcp__"
+
+_oauth_token_store = None
+
+
+def is_stdio(address: str) -> bool:
+ """A non-HTTP address is a local stdio command, e.g.
+ 'npx -y @modelcontextprotocol/server-filesystem /path'."""
+ return not address.strip().lower().startswith(("http://", "https://"))
+
+
+def parse_stdio_command(address: str) -> list[str]:
+ """Split a stdio command line into argv. Shared by route validation and the
+ transport so both agree on quoting (notably Windows backslash paths)."""
+ posix = sys.platform != "win32"
+ parts = shlex.split(address, posix = posix)
+ if not posix:
+ # posix=False keeps backslash paths intact but also keeps the surrounding
+ # quotes on a token. Strip a matched pair so the argv reaches the
+ # subprocess clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node).
+ parts = [
+ p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p
+ for p in parts
+ ]
+ return parts
+
+
+def stdio_mcp_enabled() -> bool:
+ """stdio MCP servers spawn local processes as the backend user (and bypass
+ the python/terminal sandbox), so they are only allowed when the backend
+ host is the user's own machine. The Tauri desktop app sets
+ UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see main.py); advanced localhost /
+ self-hosted users can opt in with the same variable. It stays off for
+ Colab and any network (0.0.0.0) bind."""
+ return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1"
+
+
+# Probe timeouts for discovering a server's tool list. OAuth needs minutes for
+# first-connect/expired-token browser sign-in; stdio allows for first-run
+# package download (e.g. `npx -y ...`); HTTP fails fast.
+_HTTP_PROBE_TIMEOUT = 8.0
+_OAUTH_PROBE_TIMEOUT = 305.0
+_STDIO_PROBE_TIMEOUT = 60.0
+
+
+def probe_timeout(address: str, use_oauth: bool) -> float:
+ if use_oauth:
+ return _OAUTH_PROBE_TIMEOUT
+ return _STDIO_PROBE_TIMEOUT if is_stdio(address) else _HTTP_PROBE_TIMEOUT
+
+
+def parse_server_headers(server: dict) -> Optional[dict]:
+ """Parsed headers_json. For stdio servers this dict is the process
+ environment instead of HTTP headers (see _client)."""
+ raw = server.get("headers_json")
+ if not raw:
+ return None
+ try:
+ parsed = json.loads(raw)
+ except (json.JSONDecodeError, ValueError):
+ return None
+ return parsed if isinstance(parsed, dict) else None
+
+
+def _oauth_store():
+ global _oauth_token_store
+ if _oauth_token_store is None:
+ from key_value.aio._utils.sanitization import AlwaysHashStrategy
+ from key_value.aio.stores.filetree import FileTreeStore
+ from utils.paths.storage_roots import ensure_dir, studio_root
+
+ # Hash keys/collections — fastmcp uses raw URLs like https://x.com as
+ # keys and FileTreeStore would treat the "://" as nested directories.
+ _oauth_token_store = FileTreeStore(
+ data_directory = ensure_dir(studio_root() / "mcp-oauth-tokens"),
+ key_sanitization_strategy = AlwaysHashStrategy(),
+ collection_sanitization_strategy = AlwaysHashStrategy(),
+ )
+ return _oauth_token_store
+
+
+async def clear_oauth_tokens_async(url: str) -> None:
+ """Drop any persisted OAuth tokens for ``url``. fastmcp keys tokens by
+ MCP URL, so on server delete / URL change / OAuth disable we have to
+ clear the old credentials explicitly. Otherwise re-registering the
+ same URL would silently reuse the old account's token. The entire
+ body runs inside the protected block -- store / OAuth construction
+ failing must not make the delete / update route 500."""
+ try:
+ from fastmcp.client.auth import OAuth
+
+ auth = OAuth(mcp_url = url, token_storage = _oauth_store())
+ await auth.token_storage_adapter.clear()
+ except Exception as exc: # noqa: BLE001
+ # Cleanup is best-effort; the row delete still wins.
+ logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
+
+
+def _client(url: str, headers: Optional[dict], use_oauth: bool = False):
+ from fastmcp import Client
+
+ if is_stdio(url):
+ # Belt-and-suspenders: never spawn unless stdio is enabled on this host.
+ if not stdio_mcp_enabled():
+ raise PermissionError("stdio MCP servers are disabled on this host")
+ from fastmcp.client.transports import StdioTransport
+
+ parts = parse_stdio_command(url)
+ if not parts:
+ raise ValueError(f"Empty stdio command: {url!r}")
+ # env vars ride the headers field (merged over the SDK's safe default env).
+ # keep_alive=False tears the subprocess down on exit, so a one-shot
+ # probe/tool call never leaves an orphan process.
+ return Client(
+ StdioTransport(
+ command = parts[0],
+ args = parts[1:],
+ env = headers or None,
+ keep_alive = False,
+ )
+ )
+
+ from fastmcp.client.transports import SSETransport, StreamableHttpTransport
+ from fastmcp.mcp_config import infer_transport_type_from_url
+
+ auth = None
+ if use_oauth:
+ from fastmcp.client.auth import OAuth
+
+ auth = OAuth(mcp_url = url, token_storage = _oauth_store())
+
+ transport_cls = (
+ SSETransport
+ if infer_transport_type_from_url(url) == "sse"
+ else StreamableHttpTransport
+ )
+ return Client(transport_cls(url = url, headers = headers or None, auth = auth))
+
+
+async def list_tools_async(
+ url: str,
+ headers: Optional[dict] = None,
+ timeout: float = 5.0,
+ use_oauth: bool = False,
+) -> list[dict]:
+ async def _fetch() -> list[dict]:
+ async with _client(url, headers, use_oauth) as client:
+ tools = await client.list_tools()
+ return [t.model_dump(exclude_none = True) for t in tools]
+
+ return await asyncio.wait_for(_fetch(), timeout = timeout)
+
+
+def _flatten_result(result: Any) -> str:
+ parts = []
+ for block in getattr(result, "content", None) or []:
+ text = getattr(block, "text", None)
+ if text:
+ parts.append(str(text))
+ body = "\n".join(parts)
+ if not body:
+ structured = getattr(result, "structured_content", None)
+ body = str(structured) if structured is not None else ""
+
+ if getattr(result, "is_error", False):
+ # "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge.
+ return f"Error: {body}" if body else "Error: tool returned no content"
+ return body
+
+
+def call_tool_sync(
+ url: str,
+ headers: Optional[dict],
+ name: str,
+ args: dict,
+ timeout: Optional[float] = 300.0,
+ use_oauth: bool = False,
+ cancel_event = None,
+) -> str:
+ """Synchronously call an MCP tool.
+
+ ``cancel_event``: optional ``threading.Event``. When set, the in-flight
+ HTTP call is cancelled and the function returns a cancellation Error.
+ Polled in parallel with the tool call via ``asyncio.wait`` so a /cancel
+ POST from the UI interrupts even mid-network-read.
+ """
+
+ async def _call() -> Any:
+ async with _client(url, headers, use_oauth) as client:
+ return await client.call_tool(name, args)
+
+ async def _watch_cancel() -> None:
+ # 50 ms cadence keeps cancellation responsive without busy-looping;
+ # matches the cadence routes/inference.py uses for cancel watchers.
+ while cancel_event is not None and not cancel_event.is_set():
+ await asyncio.sleep(0.05)
+
+ async def _race() -> Any:
+ # Check cancellation before spawning the call task so a pre-set
+ # event short-circuits before opening the transport / HTTP
+ # connection (reviewer-reproduced race).
+ if cancel_event is not None and cancel_event.is_set():
+ raise _MCPCancelled
+ call_task = asyncio.create_task(_call())
+ if cancel_event is None:
+ return await asyncio.wait_for(call_task, timeout = timeout)
+ watch_task = asyncio.create_task(_watch_cancel())
+ try:
+ done, pending = await asyncio.wait(
+ {call_task, watch_task},
+ timeout = timeout,
+ return_when = asyncio.FIRST_COMPLETED,
+ )
+ finally:
+ for t in (call_task, watch_task):
+ if not t.done():
+ t.cancel()
+ if not done:
+ raise asyncio.TimeoutError
+ if call_task in done:
+ return call_task.result()
+ raise _MCPCancelled
+
+ try:
+ result = asyncio.run(_race())
+ except _MCPCancelled:
+ return f"Error: MCP tool '{name}' cancelled"
+ except asyncio.TimeoutError:
+ return f"Error: MCP tool '{name}' timed out after {timeout:g}s"
+ except Exception as exc:
+ logger.exception("MCP call_tool failed for %s: %s", name, exc)
+ return f"Error: MCP tool '{name}' failed: {exc}"
+
+ return _flatten_result(result)
+
+
+class _MCPCancelled(Exception):
+ """Internal sentinel raised when cancel_event fires before the tool returns."""
diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py
index fef9ba3e12..785f1dec3b 100644
--- a/studio/backend/core/inference/providers.py
+++ b/studio/backend/core/inference/providers.py
@@ -65,28 +65,77 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
},
"gemini": {
"display_name": "Google Gemini",
- "base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
- # Curated lineup — Google's /v1beta/openai/models returns dozens
- # of historical / experimental / embedding ids. Cap to the current
- # 3.x family plus the rolling `*-latest` aliases.
+ # Native Gemini REST endpoint -- the Gemini API does NOT speak
+ # OpenAI Chat Completions on this base. Requests/responses are
+ # translated in `_stream_gemini` in external_provider.py.
+ # API reference: https://ai.google.dev/gemini-api/docs
+ "base_url": "https://generativelanguage.googleapis.com/v1beta",
+ # Curated lineup -- the live ListModels response returns dozens
+ # of historical / experimental / embedding ids. Cap to the
+ # current chat-capable Gemini families (3.5 / 3.1 / 3 Flash /
+ # 2.5) plus the Nano Banana image trio and the rolling
+ # `*-latest` aliases. Excluded on purpose:
+ # - `gemini-2.0-flash*` (Google retired 2026-06-01; 404 on use)
+ # - `gemini-3-pro-preview` (shut down 2026-03-09; auto-redirects
+ # to `gemini-3.1-pro-preview` per Google's deprecation notice,
+ # so we surface 3.1 directly and skip the redirect).
+ # The allowlist below blocks the retired ids from re-appearing
+ # via the live ListModels fetch. Verified against the live
+ # `/v1beta/models` catalog 2026-05-24.
"default_models": [
"gemini-3.1-pro-preview",
+ "gemini-3.5-flash",
"gemini-3.1-flash-lite",
"gemini-3-flash-preview",
"gemini-pro-latest",
"gemini-flash-latest",
"gemini-flash-lite-latest",
+ "gemini-2.5-pro",
+ "gemini-2.5-flash",
+ "gemini-2.5-flash-lite",
+ "gemini-3-pro-image-preview",
+ "gemini-3.1-flash-image-preview",
+ "gemini-2.5-flash-image",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
- "auth_header": "Authorization",
- "auth_prefix": "Bearer ",
- "notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.",
+ # The native API takes the API key on the `x-goog-api-key`
+ # header. An empty `auth_prefix` ensures we send the bare key.
+ "auth_header": "x-goog-api-key",
+ "auth_prefix": "",
+ "openai_compatible": False,
+ "notes": (
+ "Native Gemini API. Translation lives in _stream_gemini. "
+ "API key from https://aistudio.google.com/apikey. "
+ "See https://ai.google.dev/gemini-api/docs for endpoint shapes."
+ ),
+ # Even after the regex match, drop ids that Google still
+ # returns from ListModels but routes via implicit redirect.
+ # gemini-3-pro-preview was shut down 2026-03-09 and is
+ # auto-aliased to gemini-3.1-pro-preview; we surface the
+ # canonical id only so users do not see two cards for the
+ # same underlying model.
+ "model_id_deny_exact": ("gemini-3-pro-preview",),
+ # Matches the chat-capable 3.5 / 3.1 / 3 / 2.5 families plus the
+ # rolling *-latest aliases (which Google rolls forward as new
+ # generations ship). Image-tier ids (`-image`, `-image-preview`,
+ # `nano-banana-pro-preview`) flow through the Nano Banana
+ # `responseModalities` path in `_stream_gemini`. Retired 2.0
+ # ids ARE NOT in this regex on purpose -- Google's ListModels
+ # would otherwise re-surface them and they 404 on use.
"model_id_allowlist": re.compile(
- r"^(gemini-3\.1-flash-lite|gemini-3-flash-preview|"
- r"gemini-3\.1-pro-preview|gemini-pro-latest|"
- r"gemini-flash-latest|gemini-flash-lite-latest)$"
+ r"^("
+ r"gemini-3\.5-(?:flash|pro)(?:-preview)?|"
+ r"gemini-3\.1-(?:flash|pro|flash-lite)(?:-preview)?(?:-customtools)?|"
+ r"gemini-3\.1-flash-image-preview|"
+ r"gemini-3-(?:flash|pro)(?:-preview)?|"
+ r"gemini-3-pro-image-preview|"
+ r"nano-banana-pro-preview|"
+ r"gemini-2\.5-pro|gemini-2\.5-flash|gemini-2\.5-flash-lite|"
+ r"gemini-2\.5-flash-image|"
+ r"gemini-pro-latest|gemini-flash-latest|gemini-flash-lite-latest"
+ r")$"
),
},
"deepseek": {
diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py
index a0ab8a2a53..2f94990623 100644
--- a/studio/backend/core/inference/tool_call_parser.py
+++ b/studio/backend/core/inference/tool_call_parser.py
@@ -13,13 +13,16 @@ import re
# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing
# unclosed runs so truncated tails don't leak markup.
+# Function-name char set tracks OpenAI's ^[a-zA-Z0-9_-]{1,64}$ so MCP
+# tool names that contain a hyphen (e.g. mcp__srv__list-issues) parse
+# the same as the built-in web_search/python/terminal names.
_TOOL_CLOSED_PATS = [
re.compile(r".*? ", re.DOTALL),
- re.compile(r".*?", re.DOTALL),
+ re.compile(r".*?", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r".*$", re.DOTALL),
- re.compile(r".*$", re.DOTALL),
+ re.compile(r".*$", re.DOTALL),
]
@@ -60,10 +63,12 @@ BUDGET_EXHAUSTED_NUDGE = (
# Pre-compiled patterns reused by ``parse_tool_calls_from_text``.
_TC_JSON_START_RE = re.compile(r"\s*\{")
-_TC_FUNC_START_RE = re.compile(r"\s*")
+_TC_FUNC_START_RE = re.compile(r"\s*")
_TC_END_TAG_RE = re.compile(r" ")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$")
-_TC_PARAM_START_RE = re.compile(r"\s*")
+# Parameter names can carry hyphens too (e.g. MCP tool schemas with
+# `issue-number`, `repo-name`); using `\w+` here dropped those keys.
+_TC_PARAM_START_RE = re.compile(r"\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$")
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index 0e9cce7c3e..baf1236456 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -14,6 +14,7 @@ import signal
os.environ["UNSLOTH_IS_PRESENT"] = "1"
+import asyncio
import random
import re
import shlex
@@ -24,6 +25,17 @@ import tempfile
import threading
import urllib.request
+from core.inference.mcp_client import (
+ MCP_TOOL_PREFIX,
+ call_tool_sync,
+ is_stdio,
+ list_tools_async,
+ parse_server_headers,
+ probe_timeout,
+ stdio_mcp_enabled,
+)
+from storage import mcp_servers_db
+
from loggers import get_logger
logger = get_logger(__name__)
@@ -505,6 +517,94 @@ TERMINAL_TOOL = {
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
+# OpenAI's function.name regex: ^[a-zA-Z0-9_-]{1,64}$ -- enforced before
+# streaming starts. MCP servers can return tool names containing '.', '/',
+# spaces, etc., which the prefix scheme would forward to OpenAI verbatim
+# and 400 the whole request. Validate up front and skip with a warning.
+_OPENAI_FN_NAME_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")
+
+
+def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
+ """Convert an MCP server's tool list into OpenAI function specs."""
+ display = server.get("display_name") or server["id"]
+ specs: list[dict] = []
+ seen_names: set[str] = set()
+ for tool in mcp_tools:
+ raw_name = tool.get("name") or ""
+ if not raw_name:
+ logger.warning("Skipping MCP tool on '%s': empty name.", display)
+ continue
+ name = f"{MCP_TOOL_PREFIX}{server['id']}__{raw_name}"
+ # OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; bad chars
+ # (., /, spaces, etc.) or oversized names would 400 the whole
+ # request. Skip + warn so the rest of the tools still ship.
+ if not _OPENAI_FN_NAME_RE.fullmatch(name):
+ logger.warning(
+ "Skipping MCP tool '%s' on '%s': composed name '%s' is not "
+ "valid OpenAI function.name (regex ^[a-zA-Z0-9_-]{1,64}$).",
+ raw_name,
+ display,
+ name,
+ )
+ continue
+ # Same MCP server returning duplicate tool names would also 400
+ # OpenAI ("tools[N].function.name duplicates ..."). Drop dupes.
+ if name in seen_names:
+ logger.warning(
+ "Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display
+ )
+ continue
+ seen_names.add(name)
+ specs.append(
+ {
+ "type": "function",
+ "function": {
+ "name": name,
+ "description": f"[{display}] {tool.get('description') or ''}".strip(),
+ "parameters": tool.get("inputSchema")
+ or {"type": "object", "properties": {}},
+ },
+ }
+ )
+ return specs
+
+
+async def get_enabled_mcp_tools() -> list[dict]:
+ servers = [s for s in mcp_servers_db.list_servers() if s.get("is_enabled")]
+ # Never spawn stdio servers when stdio is disabled on this host (e.g. a DB
+ # carried over from a desktop install onto a Colab / network deployment).
+ if not stdio_mcp_enabled():
+ servers = [s for s in servers if not is_stdio(s["url"])]
+ if not servers:
+ return []
+
+ results = await asyncio.gather(
+ *(
+ list_tools_async(
+ url = s["url"],
+ headers = parse_server_headers(s),
+ timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))),
+ use_oauth = bool(s.get("use_oauth")),
+ )
+ for s in servers
+ ),
+ return_exceptions = True,
+ )
+
+ specs: list[dict] = []
+ for server, payload in zip(servers, results):
+ if isinstance(payload, BaseException):
+ logger.warning(
+ "MCP server '%s' (%s) discovery failed: %s",
+ server.get("display_name") or server["id"],
+ server.get("url"),
+ payload,
+ )
+ continue
+ specs.extend(_mcp_specs_for_server(server, payload))
+ return specs
+
+
_TIMEOUT_UNSET = object()
@@ -525,6 +625,27 @@ def execute_tool(
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
)
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
+ if name.startswith(MCP_TOOL_PREFIX):
+ try:
+ _, server_id, tool_name = name.split("__", 2)
+ except ValueError:
+ return f"Error: malformed MCP tool name '{name}'"
+ server = mcp_servers_db.get_server(server_id)
+ if not server:
+ return f"Error: MCP server '{server_id}' not found"
+ if not server.get("is_enabled"):
+ return f"Error: MCP server '{server_id}' is disabled"
+ if is_stdio(server["url"]) and not stdio_mcp_enabled():
+ return f"Error: stdio MCP server '{server_id}' is disabled on this host"
+ return call_tool_sync(
+ url = server["url"],
+ headers = parse_server_headers(server),
+ name = tool_name,
+ args = arguments,
+ timeout = effective_timeout,
+ use_oauth = bool(server.get("use_oauth")),
+ cancel_event = cancel_event,
+ )
if name == "web_search":
return _web_search(
arguments.get("query", ""),
@@ -632,8 +753,17 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
for *_, sockaddr in infos:
ip = ipaddress.ip_address(sockaddr[0])
+ # `not ip.is_global` rejects every category the denylist below
+ # also rejects PLUS shared address space (100.64.0.0/10 carrier-
+ # grade NAT) and benchmarking/documentation/exchange ranges that
+ # Python classifies with `is_private=False` and `is_global=False`
+ # (see https://docs.python.org/3/library/ipaddress.html#ipaddress.IPv4Address.is_global).
+ # The explicit predicates after it give human-readable categories
+ # in the error message, but a single non-global check is the
+ # source of truth and prevents future ranges from leaking.
if (
- ip.is_private
+ not ip.is_global
+ or ip.is_private
or ip.is_loopback
or ip.is_link_local
or ip.is_multicast
diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py
index bb61965764..e8bd7d9ea4 100644
--- a/studio/backend/core/tool_healing.py
+++ b/studio/backend/core/tool_healing.py
@@ -17,22 +17,25 @@ verifies it with AST comparison.
import json
import re
-# Pre-compiled patterns for tool XML stripping.
+# Pre-compiled patterns for tool XML stripping. Hyphen in the
+# function/parameter name char-class tracks OpenAI's allowed set so
+# MCP tool names with dashes (mcp__srv__list-issues) and parameter
+# names with dashes (`issue-number`) parse alongside the built-ins.
_TOOL_CLOSED_PATS = [
re.compile(r".*? ", re.DOTALL),
- re.compile(r".*?", re.DOTALL),
+ re.compile(r".*?", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r".*$", re.DOTALL),
- re.compile(r".*$", re.DOTALL),
+ re.compile(r".*$", re.DOTALL),
]
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"\s*\{")
-_TC_FUNC_START_RE = re.compile(r"\s*")
+_TC_FUNC_START_RE = re.compile(r"\s*")
_TC_END_TAG_RE = re.compile(r" ")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$")
-_TC_PARAM_START_RE = re.compile(r"\s*")
+_TC_PARAM_START_RE = re.compile(r"\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$")
diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index b128fb5338..0365b3ffd8 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -6,8 +6,10 @@ Unsloth Training Backend
Integrates Unsloth training capabilities with the FastAPI backend
"""
+import gc
import os
import sys
+import types
# Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -42,7 +44,10 @@ from utils.hardware import (
get_visible_gpu_count,
)
-torch._dynamo.config.recompile_limit = 64
+# recompile_limit was removed in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2).
+# Guard so training doesn't crash on RDNA2/RDNA3 with older ROCm torch wheels.
+if hasattr(torch._dynamo.config, "recompile_limit"):
+ torch._dynamo.config.recompile_limit = 64
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
@@ -417,8 +422,6 @@ class UnslothTrainer:
in sys.modules. When the next training run calls dataset.map(num_proc=N),
forked child processes inherit this stale state and deadlock.
"""
- import sys as _sys
-
# Remove cloned audio repo paths from sys.path
base_dir = os.path.dirname(os.path.abspath(__file__))
audio_paths = [
@@ -433,15 +436,15 @@ class UnslothTrainer:
removed_paths = []
for path in audio_paths:
- if path in _sys.path:
- _sys.path.remove(path)
+ if path in sys.path:
+ sys.path.remove(path)
removed_paths.append(path)
# Remove stale audio modules from sys.modules
prefixes = ("snac", "whisper", "sparktts", "outetts")
- removed_modules = [key for key in _sys.modules if key.startswith(prefixes)]
+ removed_modules = [key for key in sys.modules if key.startswith(prefixes)]
for key in removed_modules:
- del _sys.modules[key]
+ del sys.modules[key]
if removed_paths or removed_modules:
logger.info(
@@ -538,10 +541,9 @@ class UnslothTrainer:
# clear_unsloth_compiled_cache() deletes the disk cache, but the flag
# prevents re-compilation — leaving missing cache files. Reloading
# restores original class definitions so Unsloth can re-compile cleanly.
- import sys as _sys
import importlib
- for _key, _mod in list(_sys.modules.items()):
+ for _key, _mod in list(sys.modules.items()):
if "transformers.models." in _key and ".modeling_" in _key:
if hasattr(_mod, "__UNSLOTH_PATCHED__"):
try:
@@ -657,6 +659,23 @@ class UnslothTrainer:
f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
)
+ # AMD ROCm hardware without native bfloat16 (e.g. RDNA2 / gfx103x)
+ # crashes with an LLVM error at the first bf16 kernel dispatch if
+ # dtype=None lets unsloth auto-pick bf16. Force float16 there so that
+ # path is never reached. NVIDIA keeps dtype=None so unsloth's own
+ # bf16/fp16/float32 auto-detection (including FORCE_FLOAT32 models) is
+ # honored -- older NVIDIA without bf16 (T4/V100) must NOT be coerced to
+ # float16 here, which the previous unconditional branch did wrongly.
+ # Derive ROCm inline (not hardware.IS_ROCM) because that flag is unset
+ # until detect_hardware() runs, which isn't guaranteed in this subprocess.
+ _is_rocm = (
+ bool(getattr(torch.version, "hip", None))
+ or "rocm" in torch.__version__.lower()
+ )
+ _auto_dtype = (
+ torch.float16 if (_is_rocm and not is_bfloat16_supported()) else None
+ )
+
# Branch based on model type
if self._audio_type == "csm":
# CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False
@@ -666,7 +685,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
- dtype = None,
+ dtype = _auto_dtype,
auto_model = CsmForConditionalGeneration,
load_in_4bit = False,
device_map = device_map,
@@ -683,7 +702,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
- dtype = None,
+ dtype = _auto_dtype,
load_in_4bit = False,
device_map = device_map,
full_finetuning = full_finetuning,
@@ -705,7 +724,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
- dtype = None,
+ dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@@ -777,7 +796,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
- dtype = None,
+ dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@@ -791,7 +810,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastVisionModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
- dtype = None, # Auto-detect
+ dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@@ -824,7 +843,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
- dtype = None, # Auto-detect
+ dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@@ -1188,7 +1207,6 @@ class UnslothTrainer:
We patch at both instance AND class level for maximum reliability,
and strip non-TransformersKwargs params that Unsloth/PEFT inject.
"""
- import types
import torch
import torch.nn as nn
from transformers.models.csm.modeling_csm import (
@@ -1730,7 +1748,6 @@ class UnslothTrainer:
logger.info("Freeing SNAC codec model from GPU...\n")
snac_model.to("cpu")
del snac_model
- import gc
gc.collect()
torch.cuda.empty_cache()
@@ -1754,13 +1771,10 @@ class UnslothTrainer:
Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens),
format as special-token text strings for SFTTrainer with dataset_text_field="text".
"""
- import sys
import torch
import numpy as np
import torchaudio.transforms as T
- import subprocess
-
device = "cuda" if torch.cuda.is_available() else "cpu"
# The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo,
@@ -1960,7 +1974,6 @@ class UnslothTrainer:
audio_tokenizer.model.cpu()
audio_tokenizer.feature_extractor.cpu()
del audio_tokenizer
- import gc
gc.collect()
torch.cuda.empty_cache()
@@ -1989,7 +2002,6 @@ class UnslothTrainer:
OuteTTS AudioProcessor for speaker representations, PromptProcessor for
training prompts. Outputs text strings for SFTTrainer with dataset_text_field="text".
"""
- import sys
import io
import tempfile
import torch
@@ -2173,7 +2185,6 @@ class UnslothTrainer:
del whisper_model
del audio_processor
del prompt_processor
- import gc
gc.collect()
torch.cuda.empty_cache()
@@ -3057,6 +3068,14 @@ class UnslothTrainer:
logger.info("Configuring DeepSeek OCR data collator...\n")
FastVisionModel.for_training(self.model)
+ # DeepSeek OCR's (image_size, base_size, crop_mode) is a
+ # coupled preset; changing image_size alone desyncs the
+ # per-crop pixel grid from num_queries. Use Gundam.
+ if training_args.get("vision_image_size") is not None:
+ logger.info(
+ "Vision image resize ignored for DeepSeek OCR "
+ "(uses fixed Gundam preset).\n"
+ )
data_collator = DeepSeekOCRDataCollator(
tokenizer = self.tokenizer,
model = self.model,
@@ -3123,7 +3142,21 @@ class UnslothTrainer:
from unsloth.trainer import UnslothVisionDataCollator
FastVisionModel.for_training(self.model)
- data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
+ vision_image_size = training_args.get("vision_image_size")
+ if vision_image_size is None:
+ data_collator = UnslothVisionDataCollator(
+ self.model, self.tokenizer
+ )
+ else:
+ logger.info(
+ f"Vision image resize: {vision_image_size} (max dimension)\n"
+ )
+ data_collator = UnslothVisionDataCollator(
+ self.model,
+ self.tokenizer,
+ resize = vision_image_size,
+ resize_dimension = "max",
+ )
logger.info("Vision data collator configured\n")
# ========== TRAINING CONFIGURATION ==========
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index d2c2316d45..0af3349c6f 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -193,6 +193,7 @@ class TrainingBackend:
"hf_token": kwargs.get("hf_token", ""),
"load_in_4bit": kwargs.get("load_in_4bit", True),
"max_seq_length": kwargs.get("max_seq_length", 2048),
+ "vision_image_size": kwargs.get("vision_image_size"),
"hf_dataset": kwargs.get("hf_dataset", ""),
"local_datasets": kwargs.get("local_datasets"),
"local_eval_datasets": kwargs.get("local_eval_datasets"),
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index f47a6bd599..a825321597 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -21,6 +21,9 @@ import shutil
import sys
import time
import traceback
+import gc
+import re
+import types
import subprocess as _sp
from pathlib import Path
from typing import Any, Callable
@@ -70,6 +73,58 @@ _TILELANG_INSTALL_TIMEOUT_S = 600
_TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11")
_FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
+# Module-level handle so the torch.library.Library registration survives past
+# run_training_process() and is not garbage collected mid-run.
+_WINDOWS_ROCM_GROUPED_MM_LIB = None
+
+# Worker subprocesses inherit the parent env but not the parent's
+# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL
+# setup at module load so the first `import torch` can find amdhip64.dll even
+# when HIP_PATH\bin is not on the system PATH. Handles retained at module
+# scope so they are not garbage collected.
+_ROCM_DLL_HANDLES: list = []
+if sys.platform == "win32":
+
+ def _add_rocm_dll_dirs_worker() -> None:
+ _candidates: list[str] = []
+ for _var in ("HIP_PATH", "ROCM_PATH"):
+ _val = os.environ.get(_var)
+ if _val:
+ _candidates.append(os.path.join(_val, "bin"))
+ _default_root = os.path.join(
+ os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
+ )
+
+ def _ver_key(name: str) -> tuple:
+ # Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
+ parts = []
+ for chunk in name.split("."):
+ try:
+ parts.append((0, int(chunk)))
+ except ValueError:
+ parts.append((1, chunk))
+ return tuple(parts)
+
+ try:
+ if os.path.isdir(_default_root):
+ for _ver in sorted(
+ os.listdir(_default_root), key = _ver_key, reverse = True
+ ):
+ _bin = os.path.join(_default_root, _ver, "bin")
+ if os.path.isdir(_bin):
+ _candidates.append(_bin)
+ except OSError:
+ pass
+ for _d in _candidates:
+ if os.path.isdir(_d):
+ try:
+ _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
+ except (OSError, AttributeError):
+ pass
+
+ _add_rocm_dll_dirs_worker()
+ del _add_rocm_dll_dirs_worker
+
def _model_wants_causal_conv1d(model_name: str) -> bool:
name = model_name.lower()
@@ -320,11 +375,21 @@ def _install_package_wheel_first(
f"{snippet}",
)
else:
- logger.error(
- "Failed to install %s from PyPI:\n%s",
- display_name,
- result.stdout,
- )
+ if sys.platform == "win32":
+ # No prebuilt wheel and no source build toolchain on Windows --
+ # this is expected for packages like causal-conv1d. Log at
+ # info so users aren't alarmed by what looks like an error.
+ logger.info(
+ "%s is not available on Windows (no prebuilt wheel); skipping",
+ display_name,
+ )
+ logger.debug("Install output:\n%s", result.stdout)
+ else:
+ logger.error(
+ "Failed to install %s from PyPI:\n%s",
+ display_name,
+ result.stdout,
+ )
return False
if is_hip:
@@ -337,6 +402,9 @@ def _install_package_wheel_first(
def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
if not _model_wants_causal_conv1d(model_name):
return
+ if sys.platform == "win32":
+ logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping")
+ return
_install_package_wheel_first(
event_queue = event_queue,
@@ -404,6 +472,11 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
"""Install pinned FLA + fla-core with --no-deps. Returns True iff importable post-call."""
if os.getenv(_FLA_SKIP_ENV) == "1":
return False
+ if sys.platform == "win32":
+ logger.info(
+ "Skipping flash-linear-attention install: no prebuilt wheel for Windows"
+ )
+ return False
if sys.version_info < _FLA_MIN_PYTHON:
logger.info(
"Skipping flash-linear-attention install: requires Python >= %d.%d, have %s",
@@ -483,10 +556,17 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
return False
if result.returncode != 0:
- logger.warning(
- "flash-linear-attention install failed (continuing on torch fallback):\n%s",
- result.stdout,
- )
+ if sys.platform == "win32":
+ logger.info(
+ "flash-linear-attention not available on Windows (no prebuilt wheel); "
+ "continuing on torch fallback"
+ )
+ logger.debug("Install output:\n%s", result.stdout)
+ else:
+ logger.warning(
+ "flash-linear-attention install failed (continuing on torch fallback):\n%s",
+ result.stdout,
+ )
_send_status(
event_queue,
"flash-linear-attention install failed; continuing without it",
@@ -607,15 +687,61 @@ def _tilelang_importable() -> bool:
def _torch_has_hip() -> bool:
- """True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm."""
+ """True iff torch is a ROCm build.
+
+ `torch.version.hip` covers official PyTorch ROCm wheels; AMD SDK / Radeon
+ wheels can leave it unset but still encode "rocm" in `torch.__version__`.
+ """
try:
import torch as _torch
- return getattr(_torch.version, "hip", None) is not None
+ return bool(
+ getattr(_torch.version, "hip", None)
+ or "rocm" in getattr(_torch, "__version__", "").lower()
+ )
except Exception:
return False
+def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
+ """Classify a ROCm device as unified-memory (APU) or discrete.
+
+ Returns ``(gcn_arch, is_unified)`` where:
+ - ``gcn_arch`` is the canonical arch string (e.g. ``"gfx1151"``) when a
+ known attribute is present, or ``""`` when all arch attrs are absent.
+ - ``is_unified`` is ``True`` for AMD APUs with a shared GPU/system-RAM pool
+ (gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower
+ ``set_per_process_memory_fraction`` cap to leave headroom for the OS.
+
+ Classification priority:
+ 1. ``gcnArchName`` / variant spellings (stable, naming-independent).
+ 2. Device-name substring match as a last-resort fallback when all arch
+ attrs are absent (AMD SDK / Radeon wheels may not populate them):
+ - gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
+ - gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
+ ``Radeon 8050S`` (cut-down SKU)
+ """
+ gcn_arch = ""
+ for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"):
+ _v = (getattr(props, _attr, "") or "").split(":")[0].strip()
+ if _v:
+ gcn_arch = _v
+ break
+
+ if gcn_arch:
+ return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"}
+
+ # Arch attrs absent — fall back to device-name matching.
+ dev_lower = (getattr(props, "name", "") or "").lower()
+ is_unified = (
+ "890m" in dev_lower
+ or "880m" in dev_lower
+ or "8060s" in dev_lower
+ or "8050s" in dev_lower
+ )
+ return gcn_arch, is_unified
+
+
def _tilelang_platform_supported() -> bool:
"""True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch.
@@ -881,6 +1007,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
_ensure_tilelang_backend_unconditional(eq)
def _causal_conv1d_install(eq: Any) -> bool:
+ if sys.platform == "win32":
+ logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping")
+ return False
ok = _install_package_wheel_first(
event_queue = eq,
import_name = "causal_conv1d",
@@ -959,7 +1088,47 @@ def _activate_transformers_version(model_name: str) -> None:
activate_transformers_for_subprocess(model_name)
-def _adapt_for_mlx_vlm(items):
+def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]:
+ if width <= 0 or height <= 0 or target <= 0:
+ return width, height
+ largest_side = max(width, height)
+ if largest_side <= target:
+ return width, height
+ # Integer formula matches unsloth_zoo's collator (Python round() differs
+ # by 1px on half-pixel cases). max(1, _) avoids zero-side degenerate output.
+ new_w = max(1, (width * target + largest_side // 2) // largest_side)
+ new_h = max(1, (height * target + largest_side // 2) // largest_side)
+ return new_w, new_h
+
+
+def _resize_mlx_vlm_image(image, resize):
+ if resize is None:
+ return image
+ try:
+ from PIL import Image
+ import numpy as np
+ except ImportError:
+ return image
+ if not isinstance(image, Image.Image):
+ return image
+ image = image.convert("RGB")
+ new_size = _mlx_vlm_max_resized_size(*image.size, int(resize))
+ if new_size != image.size:
+ resampling = getattr(Image, "Resampling", Image).LANCZOS
+ image = image.resize(new_size, resampling)
+ # When a resize is requested, hand mlx-vlm a writable RGB ndarray so its
+ # PIL-path square-resize is skipped and HF processors don't warn on
+ # non-writable views. resize=None (Default) above keeps the original PIL.
+ return np.array(image, copy = True)
+
+
+def _resize_mlx_vlm_images(value, resize):
+ if isinstance(value, list):
+ return [_resize_mlx_vlm_image(image, resize) for image in value]
+ return _resize_mlx_vlm_image(value, resize)
+
+
+def _adapt_for_mlx_vlm(items, resize = None):
"""Adapt GPU-path VLM dataset output for mlx-vlm consumption.
The GPU path embeds PIL images inside messages content as
@@ -979,7 +1148,7 @@ def _adapt_for_mlx_vlm(items):
if isinstance(part, dict) and part.get("type") == "image":
img = part.get("image")
if img is not None:
- images.append(img)
+ images.append(_resize_mlx_vlm_image(img, resize))
new_content.append({"type": "image"})
else:
new_content.append(part)
@@ -990,9 +1159,9 @@ def _adapt_for_mlx_vlm(items):
if images:
out["image"] = images[0] if len(images) == 1 else images
elif "image" in item:
- out["image"] = item["image"]
+ out["image"] = _resize_mlx_vlm_images(item["image"], resize)
elif "images" in item:
- out["images"] = item["images"]
+ out["images"] = _resize_mlx_vlm_images(item["images"], resize)
adapted.append(out)
return adapted
@@ -1093,7 +1262,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
Mirrors the event_queue protocol so the parent process pump works unchanged.
"""
import time
- import gc
import math
import threading
import queue as _queue
@@ -1168,6 +1336,25 @@ def _run_mlx_training(event_queue, stop_queue, config):
is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False))
model._is_vlm_model = is_vlm
+ vision_image_size = config.get("vision_image_size")
+ # DeepSeek OCR uses a coupled preset tuple; skip resize like the Torch path.
+ _model_name_lower = str(config.get("model_name", "")).lower()
+ _is_deepseek_ocr = "deepseek" in _model_name_lower and "ocr" in _model_name_lower
+ if is_vlm and vision_image_size is not None and _is_deepseek_ocr:
+ _send(
+ "status",
+ status_message = (
+ "MLX vision image resize ignored for DeepSeek OCR "
+ "(uses fixed Gundam preset)."
+ ),
+ )
+ vision_image_size = None
+ elif is_vlm and vision_image_size is not None:
+ vision_image_size = int(vision_image_size)
+ _send(
+ "status",
+ status_message = f"MLX vision image resize: {vision_image_size} (max dimension)",
+ )
# ── 2. Apply LoRA / full FT ──
# Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.)
@@ -1302,7 +1489,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
progress_callback = _fmt_progress,
)
if vlm_info.get("success"):
- dataset = _adapt_for_mlx_vlm(vlm_info["dataset"])
+ dataset = _adapt_for_mlx_vlm(
+ vlm_info["dataset"],
+ resize = vision_image_size,
+ )
else:
errors = vlm_info.get("errors", [])
raise ValueError(
@@ -1317,7 +1507,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
dataset_name = hf_dataset or "local",
)
if ev_info.get("success"):
- eval_dataset = _adapt_for_mlx_vlm(ev_info["dataset"])
+ eval_dataset = _adapt_for_mlx_vlm(
+ ev_info["dataset"],
+ resize = vision_image_size,
+ )
elif format_type:
_send("status", status_message = f"Formatting dataset ({format_type})...")
@@ -1828,6 +2021,340 @@ def run_training_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
+ # ── 1d. Stub torchao on Windows ROCm ──
+ # Shared with the export worker; see core/_torchao_stub.py for the full
+ # rationale (torchao -> torch.distributed._functional_collectives crashes on
+ # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm.
+ # Must run before any import of transformers / unsloth_zoo.
+ from core._torchao_stub import install_torchao_windows_rocm_stub
+
+ install_torchao_windows_rocm_stub()
+
+ # ── 1e. Ensure torch.distributed helper attrs are present ──
+ # Single-GPU training never initialises the process group, so these helpers
+ # are never called — but transformers/trl import them unconditionally.
+ _td_stubs = {
+ "is_initialized": lambda: False,
+ "is_available": lambda: False,
+ "is_torchelastic_launched": lambda: False,
+ "get_rank": lambda: 0,
+ "get_world_size": lambda: 1,
+ "barrier": lambda: None,
+ }
+
+ try:
+ import torch.distributed as _td
+
+ for _name, _stub in _td_stubs.items():
+ if not hasattr(_td, _name):
+ setattr(_td, _name, _stub)
+ except Exception:
+ _td_mock = types.ModuleType("torch.distributed")
+ for _name, _stub in _td_stubs.items():
+ setattr(_td_mock, _name, _stub)
+ sys.modules["torch.distributed"] = _td_mock
+ try:
+ import torch as _torch
+
+ _torch.distributed = _td_mock
+ except Exception:
+ pass
+
+ # ── 1f. Windows ROCm runtime patches ──
+ # torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows),
+ # causing 0xC0000005 (access violation) during training.
+ #
+ # Root cause: the JitDecomp autograd decomposition system (NOT torch.compile)
+ # dispatches _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash.
+ # TORCHDYNAMO_DISABLE=1 stops the compiler frontend but does NOT stop
+ # JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm
+ # with a safe Python fallback.
+ #
+ # Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped
+ # (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200.
+ # We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get
+ # the real GPU kernel rather than our Python fallback.
+ #
+ # Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0.
+ #
+ # Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None,
+ # Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor
+ # offs: optional group-split offsets (MoE-style variable-size batches)
+ #
+ # torch is already in sys.modules from section 1e's `import torch.distributed`.
+ # Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past
+ # function return / mid-run GC.
+ global _WINDOWS_ROCM_GROUPED_MM_LIB
+ if sys.platform == "win32":
+ _torch_for_rocm = sys.modules.get("torch")
+ # Broad check: torch.version.hip OR "rocm" in torch.__version__.
+ # AMD SDK / Radeon Windows wheels do not always populate
+ # torch.version.hip; without the broad check the BNB version pin,
+ # dynamo-disable, and _grouped_mm fallback below silently skip
+ # (matches the torchao stub gate above and main.py).
+ _build_version_for_rocm = (
+ getattr(_torch_for_rocm, "__version__", "").lower()
+ if _torch_for_rocm is not None
+ else ""
+ )
+ _is_win_rocm_torch = bool(
+ _torch_for_rocm is not None
+ and (
+ getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
+ or "rocm" in _build_version_for_rocm
+ )
+ )
+ if _is_win_rocm_torch:
+ # Disable dynamo (belt-and-suspenders; JitDecomp patch below is the
+ # real fix, but keeping dynamo off avoids any other compile paths).
+ if "TORCHDYNAMO_DISABLE" not in os.environ:
+ os.environ["TORCHDYNAMO_DISABLE"] = "1"
+ logger.info("Windows ROCm: torch.compile (dynamo) disabled")
+
+ # BNB auto-detects the HIP version from torch.version.hip and uses
+ # it to choose which DLL to load (e.g. "7.13" → rocm713.dll).
+ # AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its
+ # version suffix does not always match the torch HIP version (e.g.
+ # torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still
+ # ships rocm72.dll). We detect the actual DLL name from the installed
+ # package and override BNB's auto-detection. "72" is a safe fallback
+ # if detection fails. Callers may override by pre-setting the var.
+ if "BNB_ROCM_VERSION" not in os.environ:
+ _bnb_rocm_ver = None
+ try:
+ import glob as _glob
+ import importlib.util as _ilu
+ import re as _re
+
+ _bnb_spec = _ilu.find_spec("bitsandbytes")
+ if _bnb_spec and _bnb_spec.submodule_search_locations:
+ _all_vers: list[str] = []
+ for _pkg_dir in _bnb_spec.submodule_search_locations:
+ for _dll in _glob.glob(
+ os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
+ ):
+ _m = _re.search(
+ r"libbitsandbytes_rocm(\d+)\.dll",
+ os.path.basename(_dll),
+ )
+ if _m:
+ _all_vers.append(_m.group(1))
+ # Pick the highest numeric suffix so that e.g. "713"
+ # wins over "72" when both variants are present.
+ # Filesystem glob order is not guaranteed, so always
+ # sort rather than stopping at the first match.
+ if _all_vers:
+ _bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
+ except Exception:
+ pass
+ _bnb_rocm_ver = _bnb_rocm_ver or "72"
+ os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
+ logger.info(
+ "Windows ROCm: set BNB_ROCM_VERSION=%s "
+ "(detected from installed BNB wheel; "
+ "overrides torch.version.hip auto-detection)",
+ _bnb_rocm_ver,
+ )
+
+ # Parse HIP version for the kernel-fix gate below.
+ # torch.version.hip can be "7.13.99004", "7.2.0", etc.
+ # AMD SDK / Radeon wheels may leave torch.version.hip unset and
+ # encode the ROCm version in torch.__version__ instead
+ # (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back
+ # to that string when version.hip is missing.
+ def _hip_ver_at_least(major: int, minor: int) -> bool:
+ _hip_str = getattr(
+ getattr(_torch_for_rocm, "version", None), "hip", None
+ )
+ if not _hip_str:
+ # Try the standard "+rocmX.Y.Z" embedded version first
+ # (e.g. "2.11.0+rocm7.13.0").
+ _ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm)
+ if _ver_match:
+ return (
+ int(_ver_match.group(1)),
+ int(_ver_match.group(2)),
+ ) >= (major, minor)
+ # AMD SDK / Radeon Windows wheels encode the build as
+ # "+rocmsdk" (e.g. "2.9.0+rocmsdk20251116") with no
+ # explicit rocmX.Y component. The rocmsdk format was
+ # introduced after the gfx120X null-kernel fix landed in
+ # ROCm 7.13, so any wheel with this suffix is new enough to
+ # have working HIP kernels. Treat as >= 7.13 rather than
+ # falling back to False and installing the Python workaround
+ # on a wheel that doesn't need it.
+ if "rocmsdk" in _build_version_for_rocm:
+ logger.debug(
+ "Windows ROCm: AMD SDK wheel detected (%r); "
+ "assuming HIP >= %d.%d (rocmsdk wheels post-date "
+ "the gfx120X null-kernel fix)",
+ _build_version_for_rocm,
+ major,
+ minor,
+ )
+ return True
+ return False
+ try:
+ _parts = [int(x) for x in str(_hip_str).split(".")[:2]]
+ if len(_parts) < 2:
+ logger.warning(
+ "Windows ROCm: torch.version.hip %r has fewer than "
+ "two components; cannot compare against %d.%d",
+ _hip_str,
+ major,
+ minor,
+ )
+ return False
+ return (_parts[0], _parts[1]) >= (major, minor)
+ except ValueError:
+ logger.warning(
+ "Windows ROCm: could not parse torch.version.hip %r as "
+ "a version number; assuming HIP < %d.%d",
+ _hip_str,
+ major,
+ minor,
+ )
+ return False
+
+ # _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12,
+ # causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+).
+ # Only install the Python fallback on the affected versions so users
+ # on 7.13+ get the real GPU kernel for MoE workloads.
+ if not _hip_ver_at_least(7, 13):
+ try:
+ import warnings as _warnings
+
+ _gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
+
+ def _grouped_mm_safe_impl(
+ self, mat2, offs = None, bias = None, out_dtype = None
+ ):
+ """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
+ _t = _torch_for_rocm
+ if offs is None:
+ # No offsets: behave like the real op, which
+ # accepts either (M, K) x (K, N) -> mm, or 3-D
+ # batched inputs -> bmm. Picking torch.mm
+ # unconditionally previously raised "self must be
+ # a matrix" on 3-D MoE workloads.
+ if self.dim() == 3 and mat2.dim() == 3:
+ result = _t.bmm(self.contiguous(), mat2.contiguous())
+ elif self.dim() == 3 and mat2.dim() == 2:
+ # Broadcast 2-D mat2 across the batch dim.
+ result = _t.matmul(self.contiguous(), mat2.contiguous())
+ elif self.dim() == 2 and mat2.dim() == 3:
+ # Broadcast 2-D self across batch via matmul semantics.
+ result = _t.matmul(self.contiguous(), mat2.contiguous())
+ else:
+ result = _t.mm(self.contiguous(), mat2.contiguous())
+ else:
+ # Grouped case: offs[i] is the exclusive end-row of
+ # group i in `self`; mat2 may be 3-D or 2-D.
+ offs_list = offs.tolist()
+ pieces = []
+ prev = 0
+ for idx, end in enumerate(offs_list):
+ end = int(end)
+ a_part = self[prev:end].contiguous()
+ if mat2.dim() == 3:
+ b_part = mat2[idx].contiguous()
+ else:
+ b_part = mat2.contiguous()
+ pieces.append(_t.mm(a_part, b_part))
+ prev = end
+ # Include any trailing rows not covered by offs
+ if prev < self.shape[0]:
+ a_tail = self[prev:].contiguous()
+ b_tail = (
+ mat2[-1].contiguous()
+ if mat2.dim() == 3
+ else mat2.contiguous()
+ )
+ pieces.append(_t.mm(a_tail, b_tail))
+ result = (
+ _t.cat(pieces, dim = 0)
+ if pieces
+ else _t.zeros(
+ 0,
+ mat2.shape[-1],
+ device = self.device,
+ dtype = self.dtype,
+ )
+ )
+ if bias is not None:
+ result = result + bias
+ if out_dtype is not None:
+ result = result.to(out_dtype)
+ elif result.dtype != self.dtype:
+ result = result.to(self.dtype)
+ return result
+
+ with _warnings.catch_warnings():
+ _warnings.simplefilter("ignore")
+ _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
+
+ _WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
+ logger.info(
+ "Windows ROCm: patched _grouped_mm CUDA dispatch "
+ "(null HIP kernel on gfx1200, ROCm ≤ 7.12 — "
+ "bypassed with Python mm fallback)"
+ )
+ except Exception as _patch_exc:
+ logger.warning(
+ "Windows ROCm: could not patch _grouped_mm — "
+ "training may crash with 0xC0000005: %s",
+ _patch_exc,
+ )
+ else:
+ logger.info(
+ "Windows ROCm: HIP >= 7.13 — _grouped_mm kernel is functional, "
+ "skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
+ )
+
+ # ── 1g. ROCm OOM guard ──
+ # On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can
+ # cause a HIP driver hang that freezes the entire system rather than
+ # raising a Python exception. set_per_process_memory_fraction caps the
+ # HIP allocator so PyTorch raises OutOfMemoryError before hitting the
+ # hardware limit, giving the UI a clean error instead of a system freeze.
+ # Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does
+ # not need this cap.
+ # Unified-memory APUs (gfx1150 Strix Point / gfx1151 Strix Halo) share GPU
+ # and system RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there.
+ # Primary classifier: gcnArchName from device properties — stable within a
+ # product family and naming-independent. AMD SDK / Radeon wheels may omit
+ # gcnArchName or expose it under a variant spelling, so we try several attr
+ # names then fall back to known device-name markers as a last resort.
+ # Non-fatal: silently skipped if torch is not importable.
+ if _hw.IS_ROCM:
+ try:
+ import torch as _torch_mem
+
+ if _torch_mem.cuda.is_available():
+ # Classify unified vs discrete via _rocm_classify_unified_memory.
+ # See that function's docstring for classification priority.
+ _props = _torch_mem.cuda.get_device_properties(0)
+ _dev_name = _props.name
+ _gcn_arch, _is_unified = _rocm_classify_unified_memory(_props)
+ if _is_unified and not _gcn_arch:
+ logger.debug(
+ "ROCm OOM guard: gcnArchName absent -- inferred "
+ "unified memory from device name %r; applying 0.80 cap",
+ _dev_name,
+ )
+ _mem_fraction = 0.80 if _is_unified else 0.90
+ _torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction)
+ logger.info(
+ "ROCm OOM guard: set_per_process_memory_fraction(%.2f) — "
+ "%s memory host (%s, %s)",
+ _mem_fraction,
+ "unified" if _is_unified else "discrete",
+ _dev_name,
+ _gcn_arch or "unknown arch",
+ )
+ except Exception as _oom_guard_err:
+ logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)
+
# ── 2. Now import ML libraries (fresh in this clean process) ──
try:
_send_status(event_queue, "Importing Unsloth...")
@@ -2248,6 +2775,7 @@ def run_training_process(
eval_dataset = eval_dataset,
eval_steps = eval_steps,
max_seq_length = config.get("max_seq_length", 2048),
+ vision_image_size = config.get("vision_image_size"),
optim = config.get("optim", "adamw_8bit"),
lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
is_cpt = is_cpt,
@@ -2281,14 +2809,38 @@ def run_training_process(
)
except Exception as exc:
- event_queue.put(
- {
- "type": "error",
- "error": str(exc),
- "stack": traceback.format_exc(limit = 20),
- "ts": time.time(),
- }
+ _exc_str = str(exc).lower()
+ _is_oom = (
+ "out of memory" in _exc_str
+ or "hip out of memory" in _exc_str
+ or "cuda out of memory" in _exc_str
+ or type(exc).__name__ == "OutOfMemoryError"
)
+ if _is_oom:
+ _oom_msg = (
+ "GPU ran out of VRAM during training.\n"
+ "To fix: reduce max_seq_length (e.g. 2048–4096), enable "
+ "gradient_checkpointing=True, lower per_device_train_batch_size, "
+ "or use a smaller model / higher quantization."
+ )
+ logger.error("Training stopped: GPU OOM — %s", exc)
+ event_queue.put(
+ {
+ "type": "error",
+ "error": _oom_msg,
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
+ else:
+ event_queue.put(
+ {
+ "type": "error",
+ "error": str(exc),
+ "stack": traceback.format_exc(limit = 20),
+ "ts": time.time(),
+ }
+ )
def _send_status(event_queue: Any, message: str) -> None:
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 689241b915..b6cf58a02c 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -12,12 +12,127 @@ from pathlib import Path as _Path
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
+# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
+# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
+# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
+if sys.platform == "win32":
+ # Retained at module scope -- os.add_dll_directory returns a handle that
+ # removes the search-path entry when garbage collected.
+ _ROCM_DLL_HANDLES: list = []
+
+ def _add_rocm_dll_dirs() -> None:
+ candidates = []
+ # 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer
+ for _var in ("HIP_PATH", "ROCM_PATH"):
+ _val = os.environ.get(_var)
+ if _val:
+ candidates.append(os.path.join(_val, "bin"))
+ # 2. Standard AMD installer location: C:\Program Files\AMD\ROCm\\bin
+ # Scan all installed versions, newest first.
+ _default_root = os.path.join(
+ os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
+ )
+
+ def _ver_key(name: str) -> tuple:
+ # Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
+ parts = []
+ for chunk in name.split("."):
+ try:
+ parts.append((0, int(chunk)))
+ except ValueError:
+ parts.append((1, chunk))
+ return tuple(parts)
+
+ try:
+ if os.path.isdir(_default_root):
+ for _ver in sorted(
+ os.listdir(_default_root), key = _ver_key, reverse = True
+ ):
+ _bin = os.path.join(_default_root, _ver, "bin")
+ if os.path.isdir(_bin):
+ candidates.append(_bin)
+ except OSError:
+ pass
+ for _d in candidates:
+ if os.path.isdir(_d):
+ try:
+ _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
+ except (OSError, AttributeError):
+ pass
+
+ _add_rocm_dll_dirs()
+ del _add_rocm_dll_dirs
+
+ # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
+ # bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm.dll
+ # where comes from torch.version.hip (e.g. "7.13..." → "713").
+ # The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without
+ # this the server process crashes with "Configured ROCm binary not found".
+ # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION
+ # before any import that pulls in bitsandbytes (mirrors worker.py logic).
+ # Gate on the rocm bnb DLL (the exact file this configures) or HIP_PATH/
+ # ROCM_PATH, not on torch.version.hip: that needed importing torch on every
+ # Windows host (NVIDIA/CPU included), adding seconds to startup. Radeon
+ # wheels without HIP_PATH still ship the rocm bnb DLL, so they are covered.
+ if "BNB_ROCM_VERSION" not in os.environ:
+ import glob as _glob
+ import logging as _logging
+
+ _hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"))
+ _bnb_rocm_ver = None
+ _found_rocm_bnb = False
+ try:
+ import importlib.util as _ilu
+
+ _bnb_spec = _ilu.find_spec("bitsandbytes")
+ # submodule_search_locations (not spec.origin) handles editable installs.
+ if _bnb_spec and _bnb_spec.submodule_search_locations:
+ import re as _re_bnb
+
+ _all_vers_main: list[str] = []
+ for _pkg_dir in _bnb_spec.submodule_search_locations:
+ for _dll in _glob.glob(
+ os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
+ ):
+ _found_rocm_bnb = True
+ _km = _re_bnb.search(
+ r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll)
+ )
+ if _km:
+ _all_vers_main.append(_km.group(1))
+ if _all_vers_main:
+ _bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
+ except Exception as _e:
+ _logging.getLogger(__name__).warning(
+ "Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'",
+ _e,
+ )
+ # rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72").
+ if _found_rocm_bnb or _hip_env:
+ _bnb_rocm_ver_final = _bnb_rocm_ver or "72"
+ os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
+ _logging.getLogger(__name__).info(
+ "Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)",
+ _bnb_rocm_ver_final,
+ )
+
# Ensure backend dir is on sys.path so _platform_compat is importable when
# main.py is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)
if _backend_dir not in sys.path:
sys.path.insert(0, _backend_dir)
+# `uvicorn main:app` bypasses run.py; seed thread caps here too.
+from utils.cpu_threads import configure_cpu_threads
+
+try:
+ configure_cpu_threads()
+except ValueError as exc:
+ _raw = os.environ.get("UNSLOTH_CPU_THREADS")
+ raise SystemExit(
+ f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}"
+ ) from None
+
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
@@ -122,6 +237,7 @@ from routes import (
export_router,
inference_router,
inference_studio_router,
+ mcp_servers_router,
models_router,
providers_router,
training_history_router,
@@ -181,6 +297,11 @@ def _load_desktop_owner() -> dict[str, str] | None:
_DESKTOP_OWNER = _load_desktop_owner()
+# The Tauri desktop app runs the backend on the owner's own machine, so local
+# stdio MCP servers are safe there. setdefault lets an explicit "0" opt out.
+if _DESKTOP_OWNER:
+ os.environ.setdefault("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
+
def _desktop_owner() -> dict[str, str] | None:
return _DESKTOP_OWNER
@@ -316,20 +437,58 @@ from starlette.requests import Request as _StarletteRequest # noqa: E402
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
+# /content is Colab's working directory — more reliable than env vars which
+# aren't always set depending on Colab runtime version.
+import importlib.util as _importlib_util
+
+_IS_COLAB = os.path.isdir("/content") and (
+ bool(os.environ.get("COLAB_BACKEND_URL"))
+ or bool(os.environ.get("COLAB_JUPYTER_IP"))
+ or _importlib_util.find_spec("google.colab") is not None
+)
+
+
def _build_csp(script_nonce: "str | None" = None) -> str:
script_src = "script-src 'self'"
if script_nonce:
script_src += f" 'nonce-{script_nonce}'"
+ # In Colab the parent frame can be colab.research.google.com, a multi-level
+ # *.prod.colab.dev subdomain (e.g. foo.region.prod.colab.dev — note: CSP
+ # wildcards only match one level, so *.prod.colab.dev misses these), or a
+ # sandboxed null-origin output iframe. Use '*' so any ancestor is allowed;
+ # Colab is already a sandboxed single-user environment.
+ frame_ancestors = "*" if _IS_COLAB else "'none'"
+
+ # In Colab the frontend is served over the Colab reverse-proxy at an HTTPS
+ # *.prod.colab.dev URL. Colab's kernel communication layer and the output
+ # iframe scaffolding inject scripts from *.prod.colab.dev and
+ # *.googleusercontent.com, and make fetch/WebSocket connections to those
+ # same origins. Widen script-src and connect-src in Colab mode so those
+ # requests are not blocked. 'unsafe-inline' for scripts is still omitted;
+ # our own inline script uses a nonce.
+ if _IS_COLAB:
+ script_src += " https://*.prod.colab.dev https://*.googleusercontent.com"
+ connect_src = (
+ "'self' blob: data: "
+ "https://huggingface.co https://datasets-server.huggingface.co "
+ "https://*.prod.colab.dev wss://*.prod.colab.dev "
+ "https://*.googleusercontent.com wss://*.googleusercontent.com"
+ )
+ else:
+ connect_src = (
+ "'self' https://huggingface.co https://datasets-server.huggingface.co"
+ )
+
return (
"default-src 'self'; "
"img-src 'self' data: blob: https://t0.gstatic.com "
"https://t1.gstatic.com https://t2.gstatic.com "
"https://t3.gstatic.com https://www.google.com; "
- "connect-src 'self' https://huggingface.co https://datasets-server.huggingface.co; "
+ f"connect-src {connect_src}; "
"style-src 'self' 'unsafe-inline'; "
f"{script_src}; "
"font-src 'self' data:; "
- "frame-ancestors 'none'; "
+ f"frame-ancestors {frame_ancestors}; "
"form-action 'self'; "
"base-uri 'self'"
)
@@ -345,7 +504,10 @@ class SecurityHeadersMiddleware(BaseHTTPMiddleware):
if nonce is not None:
del response.headers[_CSP_SCRIPT_NONCE_HEADER]
response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
- response.headers.setdefault("X-Frame-Options", "DENY")
+ # Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and
+ # DENY would block serve_kernel_port_as_iframe regardless of CSP.
+ if not _IS_COLAB:
+ response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("Referrer-Policy", "no-referrer")
response.headers.setdefault(
@@ -524,6 +686,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["
# standard /v1/chat/completions path.
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
+app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
@@ -708,8 +871,6 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes:
@font-face downloads to fail silently. Stripping the attribute
makes them regular same-origin fetches that work on any protocol.
"""
- import re as _re
-
html = html_bytes.decode("utf-8")
html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html)
return html.encode("utf-8")
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 0af9425fdc..bb1bd394d1 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -581,6 +581,14 @@ class ChatMessage(BaseModel):
None,
description = "OpenAI tool-result messages: name of the tool whose result this is.",
)
+ extra_content: Optional[dict] = Field(
+ None,
+ description = (
+ "Provider-specific extra fields the translator may read. "
+ "Gemini reads `extra_content.google.thought_signature` "
+ "from assistant messages to replay text-part signatures."
+ ),
+ )
@model_validator(mode = "after")
def _validate_role_shape(self) -> "ChatMessage":
@@ -708,6 +716,10 @@ class ChatCompletionRequest(BaseModel):
"all local tools are enabled and no server-side tools are forwarded."
),
)
+ mcp_enabled: Optional[bool] = Field(
+ None,
+ description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.",
+ )
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@@ -752,17 +764,42 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Override base URL for the external provider.",
)
- enable_prompt_caching: Optional[bool] = Field(
+ enable_prompt_caching: Optional[Union[bool, str]] = Field(
None,
description = (
"[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, "
- "attaches cache_control={type:ephemeral} to the system block so the "
- "static prefix is reused across turns. On OpenAI cloud, caching is "
- "automatic for prompts >=1024 tokens and this flag is informational. "
- "Ignored for every other provider (mistral, gemini, kimi, openrouter, "
- "vllm, local, etc.). Treated as enabled when omitted."
+ "boolean true attaches cache_control={type:ephemeral} to the system "
+ "block so the static prefix is reused across turns. On OpenAI cloud, "
+ "caching is automatic for prompts >=1024 tokens and the boolean is "
+ "informational. On Gemini, pass a string cache resource name such "
+ "as `cachedContents/abc123` to attach `cachedContent` on the native "
+ "request (boolean true is a no-op on Gemini because creating the "
+ "cache requires a separate POST /cachedContents call). Ignored for "
+ "every other provider. Treated as enabled when omitted."
),
)
+
+ @field_validator("enable_prompt_caching", mode = "before")
+ @classmethod
+ def _coerce_enable_prompt_caching(cls, value: Any) -> Any:
+ """Preserve the pre-PR coercion: the field used to be Optional[bool],
+ so callers historically sent JSON strings `"true"` / `"false"` and
+ Pydantic v1 coerced them. Widening to Optional[Union[bool, str]] for
+ Gemini cache resource names lets `"false"` slip through as a truthy
+ string. Coerce the canonical bool literals back so explicit opt-outs
+ stay opt-out."""
+ if isinstance(value, str):
+ lowered = value.strip().lower()
+ # Match Pydantic v1's BooleanField coercion table (yes/y/on/t/1
+ # and no/n/off/f/0) so opt-outs that used to parse still parse.
+ # Anything else is preserved as a string for Gemini's
+ # cachedContent resource path.
+ if lowered in ("true", "t", "1", "yes", "y", "on"):
+ return True
+ if lowered in ("false", "f", "0", "no", "n", "off"):
+ return False
+ return value
+
prompt_cache_ttl: Optional[str] = Field(
None,
description = (
diff --git a/studio/backend/models/mcp_servers.py b/studio/backend/models/mcp_servers.py
new file mode 100644
index 0000000000..c696eb0faa
--- /dev/null
+++ b/studio/backend/models/mcp_servers.py
@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+
+class McpServerCreate(BaseModel):
+ display_name: str
+ url: str
+ headers: Optional[dict[str, str]] = None
+ is_enabled: bool = True
+ use_oauth: bool = False
+
+
+class McpServerUpdate(BaseModel):
+ display_name: Optional[str] = None
+ url: Optional[str] = None
+ # Absent in request body = leave as-is; null = drop all headers; dict = set.
+ headers: Optional[dict[str, str]] = None
+ is_enabled: Optional[bool] = None
+ use_oauth: Optional[bool] = None
+
+
+class McpServerResponse(BaseModel):
+ id: str
+ display_name: str
+ url: str
+ headers: dict[str, str] = Field(default_factory = dict)
+ is_enabled: bool = True
+ use_oauth: bool = False
+ created_at: str
+ updated_at: str
+
+
+class McpServerTestRequest(BaseModel):
+ url: str
+ headers: Optional[dict[str, str]] = None
+ use_oauth: bool = False
+
+
+class McpServerProbeResult(BaseModel):
+ ok: bool
+ tool_count: int = 0
+ error: Optional[str] = None
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 7c53b0fee5..c6be1eff4e 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -5,10 +5,17 @@
Pydantic schemas for Training API
"""
+import re
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing import Any, Optional, List, Dict, Literal
+# ASCII integer with an optional single sign. Used by _check_vision_image_size
+# to reject "++512", "--256", and Unicode-digit strings ("５１２", "٥١٢") that
+# would otherwise slip through str.isdigit() + int().
+_INT_RE = re.compile(r"[+-]?[0-9]+")
+
+
_MAX_BATCH_SIZE = 4096
_MAX_GRAD_ACCUM = 4096
_MAX_STEPS = 1_000_000
@@ -18,6 +25,9 @@ _MAX_SEQ_LENGTH = 2_000_000
_MAX_LR_VALUE = 1.0
_MAX_LORA_R = 16_384
_MAX_LORA_ALPHA = 32_768
+_MIN_VISION_IMAGE_SIZE = 256
+# 2048 was the most I could get most llms to work at without getting unstable
+_MAX_VISION_IMAGE_SIZE = 2048
def _parse_lr(v: Any) -> float:
@@ -58,6 +68,10 @@ class TrainingStartRequest(BaseModel):
hf_token: Optional[str] = Field(None, description = "HuggingFace token")
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
max_seq_length: int = Field(2048, description = "Maximum sequence length")
+ vision_image_size: Optional[int] = Field(
+ None,
+ description = "Optional maximum image side length for VLM training. Null uses model default.",
+ )
trust_remote_code: bool = Field(
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
@@ -159,6 +173,40 @@ class TrainingStartRequest(BaseModel):
)
return v
+ @field_validator("vision_image_size", mode = "before")
+ @classmethod
+ def _check_vision_image_size(cls, v: Any) -> Optional[int]:
+ # mode="before" sees True/False as bool (not 1/0) for a precise error.
+ if v is None:
+ return v
+ if isinstance(v, bool):
+ raise ValueError("vision_image_size must be an integer or null")
+ if isinstance(v, int):
+ coerced = v
+ elif isinstance(v, str) and _INT_RE.fullmatch(v.strip()):
+ coerced = int(v.strip())
+ elif isinstance(v, float) and v.is_integer():
+ coerced = int(v)
+ else:
+ # numpy ints / Integral subclasses, without a hard numpy import.
+ try:
+ import numbers
+
+ if isinstance(v, numbers.Integral):
+ coerced = int(v)
+ elif isinstance(v, numbers.Real) and float(v).is_integer():
+ coerced = int(v)
+ else:
+ raise TypeError
+ except Exception:
+ raise ValueError("vision_image_size must be an integer or null")
+ if coerced < _MIN_VISION_IMAGE_SIZE or coerced > _MAX_VISION_IMAGE_SIZE:
+ raise ValueError(
+ f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, "
+ f"{_MAX_VISION_IMAGE_SIZE}] (got {coerced!r})"
+ )
+ return coerced
+
@field_validator("warmup_steps")
@classmethod
def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]:
diff --git a/studio/backend/requirements/extras.txt b/studio/backend/requirements/extras.txt
index d783975a4f..daa8982ea5 100644
--- a/studio/backend/requirements/extras.txt
+++ b/studio/backend/requirements/extras.txt
@@ -52,6 +52,5 @@ addict
easydict
einops
tabulate
-fastmcp>=3.0.2
openai>=2.7.2
websockets>=15.0.1
diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt
index 96f8816b57..d6eba73245 100644
--- a/studio/backend/requirements/studio.txt
+++ b/studio/backend/requirements/studio.txt
@@ -18,3 +18,4 @@ diceware
ddgs
cryptography>=42.0.0
httpx>=0.27.0
+fastmcp>=3.0.2
diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py
index 6bb5d15e8e..ee3ab61b6e 100644
--- a/studio/backend/routes/__init__.py
+++ b/studio/backend/routes/__init__.py
@@ -16,6 +16,7 @@ from routes.export import router as export_router
from routes.training_history import router as training_history_router
from routes.chat_history import router as chat_history_router
from routes.providers import router as providers_router
+from routes.mcp_servers import router as mcp_servers_router
__all__ = [
"training_router",
@@ -29,4 +30,5 @@ __all__ = [
"training_history_router",
"chat_history_router",
"providers_router",
+ "mcp_servers_router",
]
diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py
index 1f5c0f34e0..7184934ce9 100644
--- a/studio/backend/routes/data_recipe/mcp.py
+++ b/studio/backend/routes/data_recipe/mcp.py
@@ -36,8 +36,18 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
providers: list[McpToolsProviderResult] = []
tool_to_providers: dict[str, list[str]] = defaultdict(list)
+ from core.inference.mcp_client import stdio_mcp_enabled
+
for provider_payload in payload.mcp_providers:
provider_name = str(provider_payload.get("name", "")).strip()
+ if provider_payload.get("provider_type") == "stdio" and not stdio_mcp_enabled():
+ providers.append(
+ McpToolsProviderResult(
+ name = provider_name,
+ error = "Local (stdio) MCP servers are disabled on this host.",
+ )
+ )
+ continue
built = build_mcp_providers({"mcp_providers": [provider_payload]})
if len(built) != 1:
providers.append(
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index a156f2397c..7d1c7b2488 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -435,7 +435,9 @@ _TOOL_ACTION_NUDGE = (
# 4. tail-only `` (outer close truncated by EOS); anchored to
# `\Z` so mid-text `` in user code samples survives.
_TOOL_XML_RE = _re.compile(
- r"<(?:tool_call|function=\w+)>.*?(?:|\Z)"
+ # Hyphen in the name char-class matches MCP tool names with dashes
+ # (mcp__srv__list-issues) which would otherwise leak past this strip.
+ r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)"
r"|"
r"| \s*\Z",
_re.DOTALL,
@@ -1705,6 +1707,7 @@ def _build_external_messages(
messages: list,
supports_vision: bool,
provider_type: Optional[str] = None,
+ base_url: Optional[str] = None,
) -> list[dict]:
"""
Convert ChatMessage list to OpenAI-compatible dicts for external providers.
@@ -1732,14 +1735,171 @@ def _build_external_messages(
document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS
anthropic = provider_type == "anthropic"
openai = provider_type == "openai"
+ # `extra_content` is a Gemini-specific carrier for the assistant's
+ # text-part `thoughtSignature` round-trip on the native
+ # streamGenerateContent endpoint. Custom Gemini OpenAI-compatible
+ # gateways (LiteLLM etc.) route through /chat/completions where
+ # the field is unknown and can be rejected -- gate strictly on the
+ # Google-hosted Gemini base.
+ _native_gemini = False
+ if provider_type == "gemini" and base_url:
+ try:
+ from urllib.parse import urlparse as _urlparse
+
+ _host = (_urlparse(base_url).hostname or "").lower()
+ _native_gemini = _host == "generativelanguage.googleapis.com"
+ except Exception:
+ _native_gemini = False
+ emit_extra_content = _native_gemini
+
+ _SERVER_BUILTIN_TOOL_NAMES = frozenset(
+ {"web_search", "web_fetch", "code_execution", "image_generation"}
+ )
+
+ def _is_marked_server_builtin_tool_call(tc: Any) -> bool:
+ """Return True iff `tc` is a synthetic provider-side tool card
+ with one of the canonical builtin names and either:
+ - the new `args._server_tool` marker stamped by the backend, or
+ - a Gemini `args.google.native_part` payload (durable replay
+ signal for code_execution / image_generation that predates
+ the marker).
+ Such cards must not be forwarded to non-native providers
+ because they are not real user functions and the receiving API
+ will reject the orphan tool history. Real user functions with
+ these names normally have neither signal.
+ """
+ if not isinstance(tc, dict):
+ return False
+ fn = tc.get("function")
+ if not isinstance(fn, dict):
+ return False
+ name = (fn.get("name") or "").lower()
+ if name not in _SERVER_BUILTIN_TOOL_NAMES:
+ return False
+ raw_args = fn.get("arguments") or ""
+ try:
+ args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
+ except Exception:
+ return False
+ if not isinstance(args, dict):
+ return False
+ if args.get("_server_tool") is True:
+ return True
+ google = args.get("google")
+ return isinstance(google, dict) and isinstance(google.get("native_part"), dict)
+
+ # When we drop a server-side builtin tool_call here, the matching
+ # `role="tool"` follow-up must also be dropped from the outbound
+ # history -- otherwise the provider receives an orphan
+ # tool_call_id with no matching assistant call, which OpenAI
+ # Responses and Anthropic both reject.
+ dropped_server_builtin_tool_call_ids: set[str] = set()
+
+ def _filter_tool_calls(tool_calls: Any) -> Optional[list]:
+ """Sanitize assistant `tool_calls` for non-native-Gemini providers.
+
+ Two concerns:
+ 1. `tool_calls[i].extra_content` carries Gemini-only
+ thoughtSignature metadata; strip it for providers that
+ cannot parse the unknown key.
+ 2. Marked server-side builtin cards (`_server_tool: true` on
+ a canonical builtin name, or a Gemini `native_part`
+ payload) are provider-internal Studio tool cards from a
+ prior native Gemini turn; forwarding them to OpenAI /
+ Anthropic / custom OAI-compat gateways sends an orphan
+ `tool_calls` entry (no matching tool declaration, often
+ no matching `role="tool"` reply) that can be rejected.
+ We record the dropped call_ids so the matching role=tool
+ message is also skipped below.
+ Native Gemini keeps both untouched so the native translator can
+ replay them via `native_part`.
+ """
+ if not tool_calls:
+ return None
+ if not isinstance(tool_calls, list):
+ return tool_calls
+ if emit_extra_content:
+ return tool_calls
+ cleaned: list = []
+ for _tc in tool_calls:
+ if _is_marked_server_builtin_tool_call(_tc):
+ _tc_id = _tc.get("id") if isinstance(_tc, dict) else None
+ if isinstance(_tc_id, str) and _tc_id:
+ dropped_server_builtin_tool_call_ids.add(_tc_id)
+ continue
+ if not isinstance(_tc, dict):
+ cleaned.append(_tc)
+ continue
+ if "extra_content" not in _tc:
+ cleaned.append(_tc)
+ continue
+ _stripped = {k: v for k, v in _tc.items() if k != "extra_content"}
+ cleaned.append(_stripped)
+ return cleaned
+
result = []
for msg in messages:
+ # Drop role=tool messages whose matching server-builtin
+ # tool_call was already filtered above. Forwarding an orphan
+ # tool_result with no matching tool_call would be rejected by
+ # OpenAI Responses and Anthropic.
+ if (
+ msg.role == "tool"
+ and isinstance(msg.tool_call_id, str)
+ and msg.tool_call_id in dropped_server_builtin_tool_call_ids
+ ):
+ continue
if isinstance(msg.content, str):
- # Skip assistant messages with empty content (some providers reject them)
- if msg.role == "assistant" and not msg.content.strip():
+ # Drop bare assistant messages with no content AND no
+ # tool_calls (some providers reject empty assistant turns).
+ # Preserve assistant turns whose only payload is tool_calls
+ # so multi-turn function-call loops round-trip.
+ if (
+ msg.role == "assistant"
+ and not msg.content.strip()
+ and not msg.tool_calls
+ ):
continue
- result.append({"role": msg.role, "content": msg.content})
- elif isinstance(msg.content, list):
+ out: dict[str, Any] = {"role": msg.role, "content": msg.content}
+ if msg.role == "assistant" and msg.tool_calls:
+ _tcs = _filter_tool_calls(msg.tool_calls)
+ if _tcs:
+ out["tool_calls"] = _tcs
+ elif not msg.content.strip():
+ # Every tool_call was a synthetic provider-side
+ # card and was dropped; the assistant turn would
+ # be an empty `{"role":"assistant","content":""}`
+ # which some providers reject. Skip it entirely.
+ continue
+ if msg.role == "tool":
+ if msg.tool_call_id:
+ out["tool_call_id"] = msg.tool_call_id
+ if msg.name:
+ out["name"] = msg.name
+ if emit_extra_content and msg.role == "assistant" and msg.extra_content:
+ out["extra_content"] = msg.extra_content
+ result.append(out)
+ continue
+ # Assistant messages with content=None but populated tool_calls
+ # are valid (post-tool-call assistant turn). Forward them so the
+ # provider helper can rebuild the functionCall part.
+ if msg.content is None and msg.role == "assistant" and msg.tool_calls:
+ _filtered_tcs = _filter_tool_calls(msg.tool_calls)
+ if not _filtered_tcs:
+ # Every tool_call on this turn was provider-side
+ # synthetic and dropped; skipping the whole message
+ # avoids forwarding an empty assistant turn.
+ continue
+ _assistant_only: dict[str, Any] = {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": _filtered_tcs,
+ }
+ if emit_extra_content and msg.extra_content:
+ _assistant_only["extra_content"] = msg.extra_content
+ result.append(_assistant_only)
+ continue
+ if isinstance(msg.content, list):
if supports_vision:
parts = []
for part in msg.content:
@@ -1797,9 +1957,27 @@ def _build_external_messages(
# provider would 400 on the unknown part, so
# gate by provider_type.
parts.append({"type": "compaction", "content": part.content})
- if msg.role == "assistant" and not parts:
+ entry: dict[str, Any] = {"role": msg.role, "content": parts}
+ if msg.role == "assistant" and msg.tool_calls:
+ _tcs = _filter_tool_calls(msg.tool_calls)
+ if _tcs:
+ entry["tool_calls"] = _tcs
+ elif not parts:
+ # All tool_calls were synthetic and dropped,
+ # and no preserved content parts survived.
+ # Skip rather than forward an empty assistant
+ # turn that downstream providers reject.
+ continue
+ elif msg.role == "assistant" and not parts:
continue
- result.append({"role": msg.role, "content": parts})
+ if msg.role == "tool":
+ if msg.tool_call_id:
+ entry["tool_call_id"] = msg.tool_call_id
+ if msg.name:
+ entry["name"] = msg.name
+ if emit_extra_content and msg.role == "assistant" and msg.extra_content:
+ entry["extra_content"] = msg.extra_content
+ result.append(entry)
else:
# Non-vision provider: strip images / documents, keep
# text, optionally keep compaction (Anthropic only --
@@ -1835,9 +2013,32 @@ def _build_external_messages(
if len(preserved) == 1 and preserved[0]["type"] == "text":
# Single text part collapses back to a string for
# providers that don't accept content arrays.
- result.append({"role": msg.role, "content": preserved[0]["text"]})
+ entry = {"role": msg.role, "content": preserved[0]["text"]}
else:
- result.append({"role": msg.role, "content": preserved})
+ entry = {"role": msg.role, "content": preserved}
+ if msg.role == "assistant" and msg.tool_calls:
+ _tcs = _filter_tool_calls(msg.tool_calls)
+ if _tcs:
+ entry["tool_calls"] = _tcs
+ else:
+ # All tool_calls were synthetic and dropped;
+ # skip if there's no surviving content either.
+ _entry_content = entry.get("content")
+ _has_text = (
+ isinstance(_entry_content, str) and _entry_content.strip()
+ ) or (
+ isinstance(_entry_content, list) and len(_entry_content) > 0
+ )
+ if not _has_text:
+ continue
+ if msg.role == "tool":
+ if msg.tool_call_id:
+ entry["tool_call_id"] = msg.tool_call_id
+ if msg.name:
+ entry["name"] = msg.name
+ if emit_extra_content and msg.role == "assistant" and msg.extra_content:
+ entry["extra_content"] = msg.extra_content
+ result.append(entry)
return result
@@ -1912,6 +2113,7 @@ async def _proxy_to_external_provider(
payload.messages,
_supports_vision,
provider_type = provider_type,
+ base_url = base_url,
)
client = ExternalProviderClient(
@@ -1920,6 +2122,14 @@ async def _proxy_to_external_provider(
api_key = api_key,
)
+ # `top_k` defaults to 20 in ChatCompletionRequest because the local
+ # inference path expects an int, but the external-provider path
+ # should treat "field omitted from JSON" as "use provider default"
+ # so callers that send only model/messages do not silently get
+ # different sampling than before this PR. Pydantic's
+ # `model_fields_set` tracks explicit-vs-default per request.
+ _top_k_explicit = payload.top_k if "top_k" in payload.model_fields_set else None
+
async def _stream():
gen = client.stream_chat_completion(
messages = chat_messages,
@@ -1928,7 +2138,7 @@ async def _proxy_to_external_provider(
top_p = payload.top_p,
max_tokens = payload.max_tokens,
presence_penalty = payload.presence_penalty,
- top_k = payload.top_k,
+ top_k = _top_k_explicit,
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
enabled_tools = payload.enabled_tools,
@@ -1937,6 +2147,8 @@ async def _proxy_to_external_provider(
anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id,
prompt_cache_ttl = payload.prompt_cache_ttl,
compaction_threshold = payload.compaction_threshold,
+ tools = payload.tools,
+ tool_choice = payload.tool_choice,
fast_mode = payload.fast_mode,
stream = payload.stream,
)
@@ -2438,17 +2650,29 @@ async def openai_chat_completions(
# ── Tool-calling path (agentic loop) ──────────────────
# `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools`
# hard-override the per-request value. Without a CLI override, falls
- # back to `payload.enable_tools` (existing behavior).
+ # back to `payload.enable_tools` (existing behavior). `mcp_enabled=true`
+ # also opens the tool loop so MCP-only callers do not have to flip a
+ # second flag, BUT must still honor a CLI `--disable-tools` policy --
+ # checking the raw policy here keeps `mcp_enabled` from re-enabling
+ # tools that the operator explicitly forbade.
+ from state.tool_policy import get_tool_policy as _get_tool_policy_g
+
+ _cli_policy = _get_tool_policy_g()
+ _tools_on = _effective_enable_tools(payload)
+ _mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False
use_tools = (
- _effective_enable_tools(payload)
+ (_tools_on or _mcp_allowed)
and llama_backend.supports_tools
and not has_gguf_image
)
if use_tools:
- from core.inference.tools import ALL_TOOLS
+ from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools
- if payload.enabled_tools is not None:
+ if not _tools_on:
+ # MCP-only request: skip built-ins, leave room for MCP tools.
+ tools_to_use = []
+ elif payload.enabled_tools is not None:
tools_to_use = [
t
for t in ALL_TOOLS
@@ -2457,6 +2681,19 @@ async def openai_chat_completions(
else:
tools_to_use = ALL_TOOLS
+ if _mcp_allowed:
+ tools_to_use = tools_to_use + await get_enabled_mcp_tools()
+
+ # Skip the tool loop when no tool actually survived, so the
+ # safetensors loop's "empty = allow all" semantic cannot reach
+ # built-in tools the caller did not opt into. Existing callers
+ # who omit enabled_tools still get ALL_TOOLS here, so this
+ # only suppresses the loop when discovery + opt-in left it
+ # genuinely empty.
+ if not tools_to_use:
+ use_tools = False
+
+ if use_tools:
# ── Tool-use system prompt nudge ──────────────────────
_tool_names = {t["function"]["name"] for t in tools_to_use}
_has_web = "web_search" in _tool_names
@@ -2854,9 +3091,12 @@ async def openai_chat_completions(
else:
try:
full_text = ""
+ completion_usage = None
for token in gguf_generate():
if isinstance(token, dict):
- continue # skip metadata dict in non-streaming path
+ if token.get("type") == "metadata":
+ completion_usage = token.get("usage")
+ continue
full_text = token
response = ChatCompletion(
@@ -2869,6 +3109,15 @@ async def openai_chat_completions(
finish_reason = "stop",
)
],
+ usage = CompletionUsage(
+ prompt_tokens = (completion_usage or {}).get("prompt_tokens")
+ or 0,
+ completion_tokens = (completion_usage or {}).get(
+ "completion_tokens"
+ )
+ or 0,
+ total_tokens = (completion_usage or {}).get("total_tokens") or 0,
+ ),
)
return JSONResponse(content = response.model_dump())
@@ -2932,8 +3181,15 @@ async def openai_chat_completions(
else 25
)
+ # Match the GGUF path: mcp_enabled also opens the tool loop on its own
+ # but must still honor a CLI `--disable-tools` policy.
+ from state.tool_policy import get_tool_policy as _get_tool_policy_sf
+
+ _sf_cli_policy = _get_tool_policy_sf()
+ _sf_tools_on = _effective_enable_tools(payload)
+ _sf_mcp_allowed = bool(payload.mcp_enabled) and _sf_cli_policy is not False
_sf_use_tools = (
- _effective_enable_tools(payload)
+ (_sf_tools_on or _sf_mcp_allowed)
and _sf_features.get("supports_tools", False)
and image is None
and not _sf_is_gptoss
@@ -2941,15 +3197,27 @@ async def openai_chat_completions(
)
if _sf_use_tools:
- from core.inference.tools import ALL_TOOLS
+ from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools
- if payload.enabled_tools is not None:
+ if not _sf_tools_on:
+ _sf_tools_to_use = []
+ elif payload.enabled_tools is not None:
_sf_tools_to_use = [
t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools
]
else:
_sf_tools_to_use = ALL_TOOLS
+ if _sf_mcp_allowed:
+ _sf_tools_to_use = _sf_tools_to_use + await get_enabled_mcp_tools()
+
+ # Mirror the GGUF path: refuse to enter the tool loop when nothing
+ # survived, so a model-emitted built-in call cannot piggy-back on
+ # the empty allow-list.
+ if not _sf_tools_to_use:
+ _sf_use_tools = False
+
+ if _sf_use_tools:
_sf_tool_names = {t["function"]["name"] for t in _sf_tools_to_use}
_sf_has_web = "web_search" in _sf_tool_names
_sf_has_code = "python" in _sf_tool_names or "terminal" in _sf_tool_names
@@ -4480,7 +4748,17 @@ async def anthropic_messages(
[m.model_dump() for m in payload.messages],
payload.system,
)
- openai_messages = _drop_empty_assistant_sentinels(openai_messages)
+ # Strip synthetic provider-side builtin tool history (web_search,
+ # web_fetch, code_execution, image_generation cards tagged with
+ # _server_tool or extra_content.google.native_part) before handing
+ # off to local llama-server. The local /v1/chat/completions and
+ # GGUF passthrough builders apply the same strip; without it an
+ # Anthropic /v1/messages caller replaying a prior provider-side
+ # tool_use forwards fake builtin tool history to a backend that
+ # has no matching function declarations.
+ openai_messages = _strip_provider_synthetic_tool_history(
+ _drop_empty_assistant_sentinels(openai_messages)
+ )
# Enforce vision guard + re-encode embedded images to PNG so the
# Anthropic endpoint matches the behavior of /v1/chat/completions.
@@ -5271,6 +5549,110 @@ def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]:
return out
+_LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset(
+ {"web_search", "web_fetch", "code_execution", "image_generation"}
+)
+
+
+def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]:
+ """Drop synthetic provider-side tool_calls + matching role=tool replies
+ on the local-backend (llama-server / GGUF) dispatch path.
+
+ A Gemini chat that ran code_execution / image_generation persists the
+ server-side tool card into thread history as an assistant tool_calls
+ entry tagged with ``args._server_tool`` (or a Gemini
+ ``args.google.native_part`` payload) plus a follow-up role=tool reply.
+ When the user switches the SAME thread to a local GGUF model, those
+ synthetic tool_calls are not real user functions, llama-server has no
+ matching declaration, and Gemini-only ``extra_content`` /
+ ``native_part`` payloads are meaningless. Forward only ordinary user
+ function calls; strip the matched role=tool replies too so the
+ backend does not see an orphan tool_call_id.
+ """
+ dropped_ids: set[str] = set()
+ sanitized_assistant: list[dict] = []
+ for m in messages:
+ if m.get("role") != "assistant":
+ sanitized_assistant.append(m)
+ continue
+ tool_calls = m.get("tool_calls")
+ if not isinstance(tool_calls, list) or not tool_calls:
+ # Plain text Gemini reply: still strip message-level
+ # `extra_content` (carries `google.thought_signature` replay
+ # metadata) so a text-only Gemini turn switched to a local
+ # GGUF backend does not leak Gemini-only fields to
+ # llama-server. ChatMessage previously did not have
+ # `extra_content`, so the field was implicitly dropped --
+ # round-22 added it to ChatMessage, which is what made this
+ # leak possible.
+ if "extra_content" in m:
+ m = {k: v for k, v in m.items() if k != "extra_content"}
+ sanitized_assistant.append(m)
+ continue
+ cleaned: list[dict] = []
+ for tc in tool_calls:
+ if not isinstance(tc, dict):
+ cleaned.append(tc)
+ continue
+ fn = tc.get("function")
+ name = ""
+ if isinstance(fn, dict):
+ name = (fn.get("name") or "").lower()
+ if name in _LOCAL_SERVER_BUILTIN_TOOL_NAMES:
+ raw_args = fn.get("arguments") if isinstance(fn, dict) else None
+ args_obj: Any = None
+ if isinstance(raw_args, str):
+ try:
+ args_obj = json.loads(raw_args) if raw_args else None
+ except Exception:
+ args_obj = None
+ elif isinstance(raw_args, dict):
+ args_obj = raw_args
+ is_synthetic = False
+ if isinstance(args_obj, dict):
+ if args_obj.get("_server_tool") is True:
+ is_synthetic = True
+ google = args_obj.get("google")
+ if isinstance(google, dict) and isinstance(
+ google.get("native_part"), dict
+ ):
+ is_synthetic = True
+ if is_synthetic:
+ tc_id = tc.get("id")
+ if isinstance(tc_id, str) and tc_id:
+ dropped_ids.add(tc_id)
+ continue
+ # Strip Gemini-only `extra_content` on real user tool_calls
+ # too — llama-server has no use for it and may pass it
+ # through to the model unchanged.
+ if "extra_content" in tc:
+ tc = {k: v for k, v in tc.items() if k != "extra_content"}
+ cleaned.append(tc)
+ # Drop top-level message-level `extra_content` (Gemini
+ # thoughtSignature replay metadata) on local dispatch.
+ m_clean = {k: v for k, v in m.items() if k != "extra_content"}
+ if cleaned:
+ m_clean["tool_calls"] = cleaned
+ else:
+ m_clean.pop("tool_calls", None)
+ if not m_clean.get("content") and not m_clean.get("tool_calls"):
+ continue # assistant turn now empty, drop
+ sanitized_assistant.append(m_clean)
+
+ if not dropped_ids:
+ return sanitized_assistant
+ out: list[dict] = []
+ for m in sanitized_assistant:
+ if (
+ m.get("role") == "tool"
+ and isinstance(m.get("tool_call_id"), str)
+ and m["tool_call_id"] in dropped_ids
+ ):
+ continue
+ out.append(m)
+ return out
+
+
def _openai_messages_for_passthrough(payload) -> list[dict]:
"""Build OpenAI-format message dicts for the /v1/chat/completions
passthrough path.
@@ -5287,8 +5669,10 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
``image_url`` content part so vision + function-calling requests work
transparently.
"""
- messages = _drop_empty_assistant_sentinels(
- [m.model_dump(exclude_none = True) for m in payload.messages]
+ messages = _strip_provider_synthetic_tool_history(
+ _drop_empty_assistant_sentinels(
+ [m.model_dump(exclude_none = True) for m in payload.messages]
+ )
)
if not payload.image_base64:
@@ -5337,8 +5721,10 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict]
all per-turn ``image_url`` parts so multi-image chat history keeps each
image attached to its original turn.
"""
- messages = _drop_empty_assistant_sentinels(
- [m.model_dump(exclude_none = True) for m in payload.messages]
+ messages = _strip_provider_synthetic_tool_history(
+ _drop_empty_assistant_sentinels(
+ [m.model_dump(exclude_none = True) for m in payload.messages]
+ )
)
has_message_image = any(
isinstance(msg.get("content"), list)
diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py
new file mode 100644
index 0000000000..6c63bc20ca
--- /dev/null
+++ b/studio/backend/routes/mcp_servers.py
@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import json
+import uuid
+from urllib.parse import urlparse
+
+import structlog
+from fastapi import APIRouter, Depends, HTTPException
+
+from auth.authentication import get_current_subject
+from core.inference.mcp_client import (
+ clear_oauth_tokens_async,
+ is_stdio,
+ list_tools_async,
+ parse_server_headers,
+ parse_stdio_command,
+ probe_timeout,
+ stdio_mcp_enabled,
+)
+from models.mcp_servers import (
+ McpServerCreate,
+ McpServerProbeResult,
+ McpServerResponse,
+ McpServerTestRequest,
+ McpServerUpdate,
+)
+from storage import mcp_servers_db
+
+logger = structlog.get_logger(__name__)
+
+router = APIRouter()
+
+
+def _validate_url(url: str) -> str:
+ trimmed = (url or "").strip()
+ if not trimmed:
+ raise HTTPException(status_code = 400, detail = "url must not be empty")
+ # When stdio is enabled on this host, a non-HTTP value is a local command.
+ # Reuse this field so stdio servers ride the existing CRUD/storage with no
+ # schema change. When stdio is disabled the value falls through to the
+ # http-only validation below, so non-HTTP input is just a bad URL (400).
+ if stdio_mcp_enabled() and is_stdio(trimmed):
+ try:
+ parts = parse_stdio_command(trimmed)
+ except ValueError as exc:
+ raise HTTPException(status_code = 400, detail = f"Invalid command: {exc}")
+ if not parts or not parts[0].strip():
+ raise HTTPException(status_code = 400, detail = "command must not be empty")
+ if "://" in parts[0]:
+ # A URL-scheme first token is a mistyped URL, not a command. Reject
+ # it cleanly instead of exec-ing it (mirrors the frontend check).
+ raise HTTPException(
+ status_code = 400,
+ detail = "Enter an http(s):// URL, or a local command whose "
+ "first token is an executable (not a URL).",
+ )
+ return trimmed
+ parsed = urlparse(trimmed)
+ if parsed.scheme not in ("http", "https"):
+ raise HTTPException(
+ status_code = 400,
+ detail = "url must start with http:// or https://",
+ )
+ if not parsed.netloc:
+ raise HTTPException(status_code = 400, detail = "url is missing a host")
+ return trimmed
+
+
+def _normalize_headers(headers: dict[str, str] | None) -> dict[str, str] | None:
+ """Trim header names, drop empties, coerce values to str. None if nothing left."""
+ if not headers:
+ return None
+ out: dict[str, str] = {}
+ for raw_key, value in headers.items():
+ key = str(raw_key).strip()
+ if key:
+ out[key] = str(value)
+ return out or None
+
+
+def _row_to_response(row: dict) -> McpServerResponse:
+ return McpServerResponse(
+ id = row["id"],
+ display_name = row["display_name"],
+ url = row["url"],
+ headers = parse_server_headers(row) or {},
+ is_enabled = bool(row["is_enabled"]),
+ use_oauth = bool(row.get("use_oauth")),
+ created_at = row["created_at"],
+ updated_at = row["updated_at"],
+ )
+
+
+@router.get("/", response_model = list[McpServerResponse])
+async def list_mcp_servers(
+ current_subject: str = Depends(get_current_subject),
+):
+ return [_row_to_response(row) for row in mcp_servers_db.list_servers()]
+
+
+@router.post("/", response_model = McpServerResponse, status_code = 201)
+async def create_mcp_server(
+ payload: McpServerCreate,
+ current_subject: str = Depends(get_current_subject),
+):
+ display_name = (payload.display_name or "").strip()
+ if not display_name:
+ raise HTTPException(status_code = 400, detail = "display_name must not be empty")
+ url = _validate_url(payload.url)
+ headers = _normalize_headers(payload.headers)
+ # OAuth is HTTP-only; force it off for stdio commands so a stale flag can't
+ # push the probe onto the 305s OAuth timeout. Backend is the enforcer.
+ use_oauth = payload.use_oauth and not is_stdio(url)
+
+ server_id = uuid.uuid4().hex[:16]
+ mcp_servers_db.create_server(
+ id = server_id,
+ display_name = display_name,
+ url = url,
+ headers_json = json.dumps(headers) if headers else None,
+ is_enabled = payload.is_enabled,
+ use_oauth = use_oauth,
+ )
+ return _row_to_response(mcp_servers_db.get_server(server_id))
+
+
+def _changes_from_payload(payload: McpServerUpdate) -> dict:
+ sent = payload.model_fields_set
+ changes: dict = {}
+
+ if "display_name" in sent:
+ name = (payload.display_name or "").strip()
+ if not name:
+ raise HTTPException(
+ status_code = 400, detail = "display_name must not be empty"
+ )
+ changes["display_name"] = name
+ if "url" in sent:
+ changes["url"] = _validate_url(payload.url or "")
+ if "headers" in sent:
+ headers = _normalize_headers(payload.headers)
+ changes["headers_json"] = json.dumps(headers) if headers else None
+ if "is_enabled" in sent:
+ if payload.is_enabled is None:
+ raise HTTPException(
+ status_code = 400, detail = "is_enabled must be true or false"
+ )
+ changes["is_enabled"] = payload.is_enabled
+ if "use_oauth" in sent:
+ if payload.use_oauth is None:
+ raise HTTPException(
+ status_code = 400, detail = "use_oauth must be true or false"
+ )
+ changes["use_oauth"] = payload.use_oauth
+ # stdio is OAuth-less: drop a stale OAuth flag when switching to a command.
+ if "url" in changes and is_stdio(changes["url"]):
+ changes["use_oauth"] = False
+ return changes
+
+
+@router.put("/{server_id}", response_model = McpServerResponse)
+async def update_mcp_server(
+ server_id: str,
+ payload: McpServerUpdate,
+ current_subject: str = Depends(get_current_subject),
+):
+ old = mcp_servers_db.get_server(server_id)
+ if not old:
+ raise HTTPException(status_code = 404, detail = "MCP server not found")
+ changes = _changes_from_payload(payload)
+ if not changes:
+ raise HTTPException(status_code = 400, detail = "No fields to update")
+ # headers == HTTP headers (remote) or env vars (stdio). On a transport-type
+ # switch with no new headers, drop the old ones so env secrets are not
+ # re-sent as HTTP headers (or vice versa).
+ if (
+ "url" in changes
+ and is_stdio(changes["url"]) != is_stdio(old["url"])
+ and "headers_json" not in changes
+ ):
+ changes["headers_json"] = None
+ # Clear persisted OAuth tokens when the URL changes or OAuth is
+ # disabled; fastmcp keys tokens by URL and would otherwise let a
+ # re-pointed server silently inherit the old account's credentials.
+ if bool(old.get("use_oauth")) and (
+ ("url" in changes and changes["url"] != old["url"])
+ or changes.get("use_oauth") is False
+ ):
+ await clear_oauth_tokens_async(old["url"])
+ mcp_servers_db.update_server(server_id, changes)
+ return _row_to_response(mcp_servers_db.get_server(server_id))
+
+
+@router.delete("/{server_id}", status_code = 204)
+async def delete_mcp_server(
+ server_id: str,
+ current_subject: str = Depends(get_current_subject),
+):
+ old = mcp_servers_db.get_server(server_id)
+ if not old:
+ raise HTTPException(status_code = 404, detail = "MCP server not found")
+ if old.get("use_oauth"):
+ await clear_oauth_tokens_async(old["url"])
+ mcp_servers_db.delete_server(server_id)
+
+
+@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)
+async def refresh_mcp_server_tools(
+ server_id: str,
+ current_subject: str = Depends(get_current_subject),
+):
+ server = mcp_servers_db.get_server(server_id)
+ if not server:
+ raise HTTPException(status_code = 404, detail = "MCP server not found")
+ # Refresh uses the stored address, so re-check the stdio gate here too: a
+ # stdio row from a desktop DB must not spawn on a hosted/network host.
+ if is_stdio(server["url"]) and not stdio_mcp_enabled():
+ raise HTTPException(
+ status_code = 400, detail = "stdio MCP servers are disabled on this host"
+ )
+
+ use_oauth = bool(server.get("use_oauth"))
+ try:
+ tools = await list_tools_async(
+ url = server["url"],
+ headers = parse_server_headers(server),
+ timeout = probe_timeout(server["url"], use_oauth),
+ use_oauth = use_oauth,
+ )
+ except Exception as exc: # noqa: BLE001 — surface transport+timeout errors to UI
+ logger.warning("MCP refresh failed", server_id = server_id, error = str(exc))
+ return McpServerProbeResult(ok = False, error = str(exc))
+
+ return McpServerProbeResult(ok = True, tool_count = len(tools))
+
+
+@router.post("/test", response_model = McpServerProbeResult)
+async def test_mcp_server(
+ payload: McpServerTestRequest,
+ current_subject: str = Depends(get_current_subject),
+):
+ # URL/header validation must surface as 400 like create/update so the
+ # frontend's create-form pre-flight gets the same error semantics as
+ # the actual save call. Only catch transport/timeout errors below.
+ url = _validate_url(payload.url)
+ headers = _normalize_headers(payload.headers)
+ try:
+ tools = await list_tools_async(
+ url = url,
+ headers = headers,
+ timeout = probe_timeout(url, payload.use_oauth),
+ use_oauth = payload.use_oauth,
+ )
+ except Exception as exc: # noqa: BLE001
+ return McpServerProbeResult(ok = False, error = str(exc))
+
+ return McpServerProbeResult(ok = True, tool_count = len(tools))
diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py
index 2bb1de5366..5d4bd46e62 100644
--- a/studio/backend/routes/providers.py
+++ b/studio/backend/routes/providers.py
@@ -318,22 +318,45 @@ async def list_provider_models(
try:
models = await client.list_models()
- allow_prefixes = info.get("model_id_allow_prefixes")
- if allow_prefixes is not None:
- prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p))
- if prefix_tuple:
- models = [m for m in models if m.get("id", "").startswith(prefix_tuple)]
- allowlist = info.get("model_id_allowlist")
- if allowlist is not None:
- models = [m for m in models if allowlist.match(m.get("id", ""))]
- deny_exact = info.get("model_id_deny_exact")
- if deny_exact is not None:
- deny_ids = {str(m) for m in deny_exact if str(m)}
- if deny_ids:
- models = [m for m in models if m.get("id", "") not in deny_ids]
- denylist = info.get("model_id_denylist")
- if denylist is not None:
- models = [m for m in models if not denylist.search(m.get("id", ""))]
+ # Registry-level model-id filters are scoped to the canonical
+ # native Gemini base. A custom Gemini OAI-compatible proxy
+ # (LiteLLM, deployment gateway) returns IDs like
+ # `google/gemini-2.5-flash`, `gemini/gemini-2.5-flash`, or
+ # team-prefixed deployment aliases; the native allowlist regex
+ # would strip those out and leave the picker empty even though
+ # the chat path now routes them via the OAI-compatible
+ # dispatcher (the same gate ExternalProviderClient applies for
+ # request building). Match the host check here so the model
+ # list and chat dispatch agree on what counts as "native".
+ apply_registry_model_filters = True
+ if payload.provider_type == "gemini":
+ try:
+ from urllib.parse import urlparse as _urlparse
+
+ _host = (_urlparse(base_url).hostname or "").lower()
+ except Exception:
+ _host = ""
+ apply_registry_model_filters = _host == "generativelanguage.googleapis.com"
+
+ if apply_registry_model_filters:
+ allow_prefixes = info.get("model_id_allow_prefixes")
+ if allow_prefixes is not None:
+ prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p))
+ if prefix_tuple:
+ models = [
+ m for m in models if m.get("id", "").startswith(prefix_tuple)
+ ]
+ allowlist = info.get("model_id_allowlist")
+ if allowlist is not None:
+ models = [m for m in models if allowlist.match(m.get("id", ""))]
+ deny_exact = info.get("model_id_deny_exact")
+ if deny_exact is not None:
+ deny_ids = {str(m) for m in deny_exact if str(m)}
+ if deny_ids:
+ models = [m for m in models if m.get("id", "") not in deny_ids]
+ denylist = info.get("model_id_denylist")
+ if denylist is not None:
+ models = [m for m in models if not denylist.search(m.get("id", ""))]
# Apply an optional cap after filtering so registry entries with a
# large remote catalog (e.g. HF Inference Providers) can stay
# picker-sized. No popularity sort happens server-side, so this is
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 6e2413b3e9..41a9e15562 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -194,6 +194,7 @@ async def start_training(
"hf_token": request.hf_token or "",
"load_in_4bit": request.load_in_4bit,
"max_seq_length": request.max_seq_length,
+ "vision_image_size": request.vision_image_size,
"hf_dataset": request.hf_dataset or "",
"local_datasets": request.local_datasets,
"local_eval_datasets": request.local_eval_datasets,
diff --git a/studio/backend/run.py b/studio/backend/run.py
index 3bde8abd3c..f401bc2ec1 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -19,6 +19,16 @@ backend_dir = Path(__file__).parent
if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
+from utils.cpu_threads import configure_cpu_threads
+
+try:
+ configure_cpu_threads()
+except ValueError as exc:
+ configured = os.environ.get("UNSLOTH_CPU_THREADS")
+ raise SystemExit(
+ f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}"
+ ) from None
+
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
@@ -643,7 +653,7 @@ def run_server(
from threading import Thread, Event
import uvicorn
- from main import app, setup_frontend
+ from main import app, setup_frontend, _IS_COLAB
from utils.paths import ensure_studio_directories
# Create all standard directories on startup
@@ -727,14 +737,22 @@ def run_server(
ready_event.set()
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
- config = uvicorn.Config(
- app,
+ config_kwargs = dict(
host = host,
port = port,
log_level = "info",
access_log = False,
server_header = False,
)
+ # Only in Colab: trust X-Forwarded-* from Colab's reverse proxy so the app
+ # sees the real https origin. forwarded_allow_ips="*" is fine inside Colab's
+ # single-user sandbox, but would be an unwanted security relaxation for a
+ # normal local/standalone Studio, so leave uvicorn's safe defaults
+ # (forwarded headers trusted from loopback only) in place there.
+ if _IS_COLAB:
+ config_kwargs["proxy_headers"] = True
+ config_kwargs["forwarded_allow_ips"] = "*"
+ config = uvicorn.Config(app, **config_kwargs)
_server = _ReadyServer(config)
_shutdown_event = Event()
@@ -756,14 +774,21 @@ def run_server(
app.state.trigger_shutdown = _trigger_shutdown
- # Run server in a daemon thread
+ # Run server in a daemon thread.
+ # Use an explicit new_event_loop() + run_until_complete() instead of
+ # asyncio.run() to avoid nest_asyncio's global patches to asyncio.run
+ # interfering when called from a thread while Colab/IPython already has
+ # a running loop on the main thread.
def _run():
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
try:
- asyncio.run(_server.serve())
+ loop.run_until_complete(_server.serve())
except BaseException as exc:
startup_errors.append(exc)
startup_failed.set()
finally:
+ loop.close()
if not ready_event.is_set():
startup_failed.set()
@@ -846,11 +871,33 @@ if __name__ == "__main__":
action = "store_true",
help = "API server only, no frontend (for Tauri)",
)
+ # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1
+ # applies only to direct backend launches; `unsloth studio run`
+ # always passes its own value (4) explicitly.
+ _PARALLEL_MIN = 1
+ _PARALLEL_MAX = 64
+ _PARALLEL_DEFAULT_PLAIN = 1
+ parser.add_argument(
+ "--parallel",
+ "--n-parallel",
+ type = int,
+ default = _PARALLEL_DEFAULT_PLAIN,
+ help = (
+ f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). "
+ f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4."
+ ),
+ )
args = parser.parse_args()
+ if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
+ parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
kwargs = dict(
- host = args.host, port = args.port, silent = args.silent, api_only = args.api_only
+ host = args.host,
+ port = args.port,
+ silent = args.silent,
+ api_only = args.api_only,
+ llama_parallel_slots = args.parallel,
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)
diff --git a/studio/backend/storage/mcp_servers_db.py b/studio/backend/storage/mcp_servers_db.py
new file mode 100644
index 0000000000..da2fa15423
--- /dev/null
+++ b/studio/backend/storage/mcp_servers_db.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import sqlite3
+import threading
+from datetime import datetime, timezone
+from typing import Optional
+
+from utils.paths import studio_db_path, ensure_dir
+
+_schema_lock = threading.Lock()
+_schema_ready = False
+
+
+def _ensure_schema(conn: sqlite3.Connection) -> None:
+ conn.execute("PRAGMA journal_mode=WAL")
+ conn.execute(
+ """
+ CREATE TABLE IF NOT EXISTS mcp_servers (
+ id TEXT NOT NULL PRIMARY KEY,
+ display_name TEXT NOT NULL,
+ url TEXT NOT NULL,
+ headers_json TEXT,
+ is_enabled INTEGER NOT NULL DEFAULT 1,
+ use_oauth INTEGER NOT NULL DEFAULT 0,
+ created_at TEXT NOT NULL,
+ updated_at TEXT NOT NULL
+ )
+ """
+ )
+ # use_oauth was added after the first release; backfill for pre-existing DBs.
+ cols = {
+ r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()
+ }
+ if "use_oauth" not in cols:
+ conn.execute(
+ "ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0"
+ )
+
+
+def get_connection() -> sqlite3.Connection:
+ global _schema_ready
+ db_path = studio_db_path()
+ ensure_dir(db_path.parent)
+ conn = sqlite3.connect(str(db_path))
+ conn.row_factory = sqlite3.Row
+ if not _schema_ready:
+ with _schema_lock:
+ if not _schema_ready:
+ try:
+ _ensure_schema(conn)
+ _schema_ready = True
+ except Exception:
+ conn.close()
+ raise
+ return conn
+
+
+def create_server(
+ id: str,
+ display_name: str,
+ url: str,
+ headers_json: Optional[str] = None,
+ is_enabled: bool = True,
+ use_oauth: bool = False,
+) -> None:
+ now = datetime.now(timezone.utc).isoformat()
+ conn = get_connection()
+ try:
+ conn.execute(
+ """
+ INSERT INTO mcp_servers
+ (id, display_name, url, headers_json,
+ is_enabled, use_oauth, created_at, updated_at)
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+ """,
+ (
+ id,
+ display_name,
+ url,
+ headers_json,
+ int(is_enabled),
+ int(use_oauth),
+ now,
+ now,
+ ),
+ )
+ conn.commit()
+ finally:
+ conn.close()
+
+
+def update_server(id: str, changes: dict) -> bool:
+ """Apply column updates and bump ``updated_at``. Returns True on a hit."""
+ if not changes:
+ return False
+ bool_cols = {"is_enabled", "use_oauth"}
+ sets, params = [], []
+ for col, value in changes.items():
+ sets.append(f"{col} = ?")
+ params.append(int(value) if col in bool_cols else value)
+ sets.append("updated_at = ?")
+ params.extend([datetime.now(timezone.utc).isoformat(), id])
+
+ conn = get_connection()
+ try:
+ cursor = conn.execute(
+ f"UPDATE mcp_servers SET {', '.join(sets)} WHERE id = ?",
+ params,
+ )
+ conn.commit()
+ return cursor.rowcount > 0
+ finally:
+ conn.close()
+
+
+def delete_server(id: str) -> bool:
+ conn = get_connection()
+ try:
+ cursor = conn.execute("DELETE FROM mcp_servers WHERE id = ?", (id,))
+ conn.commit()
+ return cursor.rowcount > 0
+ finally:
+ conn.close()
+
+
+def get_server(id: str) -> Optional[dict]:
+ conn = get_connection()
+ try:
+ row = conn.execute("SELECT * FROM mcp_servers WHERE id = ?", (id,)).fetchone()
+ return dict(row) if row else None
+ finally:
+ conn.close()
+
+
+def list_servers() -> list[dict]:
+ conn = get_connection()
+ try:
+ rows = conn.execute("SELECT * FROM mcp_servers ORDER BY created_at").fetchall()
+ return [dict(row) for row in rows]
+ finally:
+ conn.close()
diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py
new file mode 100644
index 0000000000..e0b819d54b
--- /dev/null
+++ b/studio/backend/tests/test_amd_apu_unified_memory.py
@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs
+(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS."""
+
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+from core.inference.llama_cpp import LlamaCppBackend
+
+
+def _fake_torch(hip, archs, *, cuda_ok = True):
+ t = types.ModuleType("torch")
+ t.version = types.SimpleNamespace(hip = hip)
+ t.cuda = types.SimpleNamespace(
+ is_available = lambda: cuda_ok,
+ device_count = lambda: len(archs),
+ get_device_properties = lambda i: types.SimpleNamespace(gcnArchName = archs[i]),
+ )
+ return t
+
+
+@pytest.mark.parametrize(
+ "hip,archs,expected",
+ [
+ ("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped)
+ ("6.2.0", ["gfx1150"], True), # Strix Point APU
+ ("6.2.0", ["gfx1100"], False), # discrete RDNA3
+ ("6.2.0", ["gfx1201"], False), # discrete RDNA4
+ ("6.2.0", ["gfx942"], False), # MI300X (data center)
+ (None, ["sm_90"], False), # NVIDIA (no torch.version.hip)
+ ("6.2.0", ["gfx1100", "gfx1151"], True), # mixed dGPU + APU
+ ],
+)
+def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected):
+ monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip, archs))
+ assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected
+
+
+def test_cpu_no_cuda_returns_false(monkeypatch):
+ monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False))
+ assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
+
+
+def test_missing_torch_returns_false(monkeypatch):
+ monkeypatch.setitem(sys.modules, "torch", None)
+ assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
diff --git a/studio/backend/tests/test_anthropic_code_execution.py b/studio/backend/tests/test_anthropic_code_execution.py
index 7f6fe58329..5c88437d17 100644
--- a/studio/backend/tests/test_anthropic_code_execution.py
+++ b/studio/backend/tests/test_anthropic_code_execution.py
@@ -275,7 +275,13 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
assert start["type"] == "tool_start"
assert start["tool_name"] == "code_execution"
assert start["tool_call_id"] == "srvtoolu_1"
- assert start["arguments"] == {"kind": "bash", "command": "ls -la"}
+ # `_server_tool: True` marks this as a provider-side synthetic
+ # tool card for the frontend's history serializer.
+ assert start["arguments"] == {
+ "kind": "bash",
+ "command": "ls -la",
+ "_server_tool": True,
+ }
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_1"
diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py
index da10d679eb..88a922e7eb 100644
--- a/studio/backend/tests/test_anthropic_web_fetch.py
+++ b/studio/backend/tests/test_anthropic_web_fetch.py
@@ -271,7 +271,12 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch):
assert start["type"] == "tool_start"
assert start["tool_name"] == "web_fetch"
assert start["tool_call_id"] == "srvtoolu_wf1"
- assert start["arguments"] == {"url": "https://example.com/article"}
+ # `_server_tool: True` marks this as a provider-side synthetic
+ # tool card for the frontend's history serializer.
+ assert start["arguments"] == {
+ "url": "https://example.com/article",
+ "_server_tool": True,
+ }
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_wf1"
# The source pill uses Title / URL / snippet as parseSourcesFromResult expects.
diff --git a/studio/backend/tests/test_cpu_threads.py b/studio/backend/tests/test_cpu_threads.py
new file mode 100644
index 0000000000..1224941622
--- /dev/null
+++ b/studio/backend/tests/test_cpu_threads.py
@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for Studio's early CPU thread-pool configuration."""
+
+import ast
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from utils.cpu_threads import _THREAD_POOL_ENV_VARS, configure_cpu_threads
+
+
+_BACKEND_DIR = Path(__file__).resolve().parent.parent
+_RUN_PY = _BACKEND_DIR / "run.py"
+_MAIN_PY = _BACKEND_DIR / "main.py"
+
+
+# Explicit positive integers seed all four native pool env vars.
+def test_cpu_thread_cap_seeds_native_pool_limits():
+ env = {"UNSLOTH_CPU_THREADS": " 6 "}
+
+ configure_cpu_threads(env)
+
+ assert {variable: env[variable] for variable in _THREAD_POOL_ENV_VARS} == {
+ variable: "6" for variable in _THREAD_POOL_ENV_VARS
+ }
+
+
+# Explicit per-library values win over the Studio knob via setdefault.
+def test_cpu_thread_cap_preserves_runtime_specific_override():
+ env = {"UNSLOTH_CPU_THREADS": "4", "OMP_NUM_THREADS": "2"}
+
+ configure_cpu_threads(env)
+
+ assert env["OMP_NUM_THREADS"] == "2"
+ assert env["MKL_NUM_THREADS"] == "4"
+
+
+# Whitespace / plus-prefix / leading zero all normalise via int().
+@pytest.mark.parametrize("raw", ["+4", "007", " 4 "])
+def test_cpu_thread_cap_normalises_valid_inputs(raw):
+ env = {"UNSLOTH_CPU_THREADS": raw}
+
+ configure_cpu_threads(env)
+
+ assert env["OMP_NUM_THREADS"] == str(int(raw.strip()))
+
+
+# Unset / empty / whitespace -> no env mutation (pure opt-in).
+@pytest.mark.parametrize("raw", [None, "", " ", "\t"])
+def test_cpu_thread_cap_is_opt_in(raw):
+ env = {} if raw is None else {"UNSLOTH_CPU_THREADS": raw}
+ snapshot = dict(env)
+
+ configure_cpu_threads(env)
+
+ assert env == snapshot
+ assert all(variable not in env for variable in _THREAD_POOL_ENV_VARS)
+
+
+# Anything that is not a positive integer raises a clear ValueError.
+@pytest.mark.parametrize(
+ "raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]
+)
+def test_cpu_thread_cap_requires_positive_integer(raw):
+ with pytest.raises(ValueError, match = "must be a positive integer"):
+ configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw})
+
+
+# env=None path uses real os.environ (production call from run.py / main.py).
+def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch):
+ for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"):
+ monkeypatch.delenv(variable, raising = False)
+ monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3")
+
+ configure_cpu_threads()
+
+ for variable in _THREAD_POOL_ENV_VARS:
+ assert os.environ[variable] == "3"
+
+
+# Calling twice must not flip any seeded value.
+def test_cpu_thread_cap_idempotent(monkeypatch):
+ for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"):
+ monkeypatch.delenv(variable, raising = False)
+ monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5")
+
+ configure_cpu_threads()
+ snapshot = {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS}
+ configure_cpu_threads()
+
+ assert {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS} == snapshot
+
+
+def _ast_line_of_configure_call(source: str) -> int:
+ tree = ast.parse(source)
+ for node in ast.walk(tree):
+ if (
+ isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Name)
+ and node.func.id == "configure_cpu_threads"
+ ):
+ return node.lineno
+ raise AssertionError("configure_cpu_threads() call not found")
+
+
+def _ast_line_of_platform_compat_import(source: str) -> int:
+ tree = ast.parse(source)
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Import):
+ for alias in node.names:
+ if alias.name == "_platform_compat":
+ return node.lineno
+ raise AssertionError("_platform_compat import not found")
+
+
+# AST-based ordering: configure_cpu_threads() must precede _platform_compat
+# in both run.py and main.py. Robust to formatting / line shifts.
+@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY])
+def test_cpu_thread_configuration_runs_before_backend_imports(entry_point):
+ source = entry_point.read_text()
+ call_line = _ast_line_of_configure_call(source)
+ compat_line = _ast_line_of_platform_compat_import(source)
+ assert call_line < compat_line, (
+ f"{entry_point.name}: configure_cpu_threads() (line {call_line}) "
+ f"must precede import _platform_compat (line {compat_line})"
+ )
+
+
+# Invalid env -> exit 1, one-line stderr, no traceback, gated before any
+# heavy import. Parametrised over both entry points.
+@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY])
+def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point):
+ env = os.environ.copy()
+ env["UNSLOTH_CPU_THREADS"] = "not-a-count"
+
+ result = subprocess.run(
+ [sys.executable, str(entry_point)],
+ env = env,
+ capture_output = True,
+ text = True,
+ )
+
+ assert result.returncode == 1
+ assert (
+ "Error: Invalid UNSLOTH_CPU_THREADS value 'not-a-count': "
+ "UNSLOTH_CPU_THREADS must be a positive integer"
+ ) in result.stderr
+ assert "Traceback" not in result.stderr
+ assert "_platform_compat" not in result.stderr
diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py
index ab1a03eeda..b1522dd382 100644
--- a/studio/backend/tests/test_desktop_auth.py
+++ b/studio/backend/tests/test_desktop_auth.py
@@ -437,6 +437,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
export_router = APIRouter(),
inference_router = APIRouter(),
inference_studio_router = APIRouter(),
+ mcp_servers_router = APIRouter(),
models_router = APIRouter(),
providers_router = APIRouter(),
training_history_router = APIRouter(),
@@ -484,11 +485,14 @@ from typer.testing import CliRunner
studio_home = Path(sys.argv[1])
real_import = builtins.__import__
-def guarded_import(name, *args, **kwargs):
+def guarded_import(name, globals = None, locals = None, fromlist = (), level = 0):
+ # Only gate absolute imports; relative `from .utils import x` inside
+ # third-party packages (e.g. typer._click.decorators) hits level > 0
+ # with name="utils" and must pass through.
blocked = ("auth", "fastapi", "structlog", "utils")
- if name in blocked or name.startswith(("auth.", "utils.")):
+ if level == 0 and (name in blocked or name.startswith(("auth.", "utils."))):
raise ModuleNotFoundError(name)
- return real_import(name, *args, **kwargs)
+ return real_import(name, globals, locals, fromlist, level)
builtins.__import__ = guarded_import
from unsloth_cli.commands import studio as studio_cli
diff --git a/studio/backend/tests/test_gemini_provider.py b/studio/backend/tests/test_gemini_provider.py
new file mode 100644
index 0000000000..4dc97302e2
--- /dev/null
+++ b/studio/backend/tests/test_gemini_provider.py
@@ -0,0 +1,5501 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Unit tests for the native Gemini API translation layer.
+
+Gemini does NOT speak OpenAI Chat Completions on its primary endpoint
+(`streamGenerateContent`). `_stream_gemini` in
+`core/inference/external_provider.py` translates between the two shapes:
+
+ Request:
+ OpenAI messages [{role, content}]
+ -> Gemini contents [{role, parts: [{text}|{inlineData}|{functionCall}|...]}]
+ + systemInstruction.parts[].text for role=system messages
+ + generationConfig.{temperature,topP,topK,maxOutputTokens}
+ + tools[{googleSearch:{}}] for web_search
+ + tools[{codeExecution:{}}] for code_execution
+ + responseModalities=[TEXT,IMAGE] for Nano Banana (gemini-2.5-flash-image)
+ + cachedContent for prompt caching
+
+ Response:
+ Gemini SSE chunks { candidates:[{content:{parts:[...]}, finishReason}],
+ usageMetadata:{promptTokenCount, candidatesTokenCount} }
+ -> OpenAI chat.completion.chunk frames
+ (delta.content for text, delta.tool_calls for functionCall,
+ _toolEvent for image_b64/web_search, usage block before [DONE])
+
+These tests pin the outbound body shape AND the inbound translation
+using httpx.MockTransport (no live network). Mirrors the structure of
+test_anthropic_cache_ttl.py and test_openai_image_generation.py.
+"""
+
+import asyncio
+import base64
+import json
+
+import httpx
+import pytest
+
+from core.inference import external_provider as ep_mod
+from core.inference.external_provider import ExternalProviderClient
+
+
+_active_mock_clients: list[httpx.AsyncClient] = []
+
+
+def _drive(coro):
+ # Create a fresh loop per drive so tests don't share asyncio state.
+ # Close mocked clients + shutdown async-generators inside this loop
+ # so Python 3.13 doesn't emit the
+ # `Response.aiter_*.aclose was never awaited` warning on GC.
+ loop = asyncio.new_event_loop()
+ try:
+ result = loop.run_until_complete(coro)
+ while _active_mock_clients:
+ mc = _active_mock_clients.pop()
+ loop.run_until_complete(mc.aclose())
+ return result
+ finally:
+ try:
+ loop.run_until_complete(loop.shutdown_asyncgens())
+ finally:
+ loop.close()
+
+
+def _make_gemini_client(
+ base_url: str = "https://generativelanguage.googleapis.com/v1beta",
+) -> ExternalProviderClient:
+ return ExternalProviderClient(
+ provider_type = "gemini",
+ base_url = base_url,
+ api_key = "AIza-test-key",
+ )
+
+
+def _mock_http(monkeypatch, handler):
+ mock_client = httpx.AsyncClient(transport = httpx.MockTransport(handler))
+ monkeypatch.setattr(ep_mod, "_http_client", mock_client)
+ # `_drive` will aclose this at the end of the run inside the same
+ # event loop so we do not leak an unawaited aclose() coroutine.
+ _active_mock_clients.append(mock_client)
+
+
+def _gemini_sse(events: list[dict]) -> bytes:
+ """Encode a list of dicts as Gemini-style SSE frames (`data:` lines)."""
+ chunks: list[str] = []
+ for event in events:
+ chunks.append(f"data: {json.dumps(event)}")
+ chunks.append("")
+ return ("\n".join(chunks) + "\n").encode("utf-8")
+
+
+def _capture_body(monkeypatch, **kwargs) -> dict:
+ """Drive a single stream and return the captured outbound request body."""
+ captured: dict = {}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ captured["headers"] = dict(request.headers)
+ captured["url"] = str(request.url)
+ captured["method"] = request.method
+ # Minimal valid Gemini stream so the helper can complete.
+ return httpx.Response(
+ 200,
+ content = _gemini_sse(
+ [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "ok"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 1,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ ),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}])
+ model = kwargs.pop("model", "gemini-2.5-flash")
+ temperature = kwargs.pop("temperature", 0.7)
+ top_p = kwargs.pop("top_p", 0.95)
+ max_tokens = kwargs.pop("max_tokens", 64)
+
+ async def run():
+ client = _make_gemini_client()
+ async for _ in client.stream_chat_completion(
+ messages = messages,
+ model = model,
+ temperature = temperature,
+ top_p = top_p,
+ max_tokens = max_tokens,
+ **kwargs,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ return captured
+
+
+def _collect(monkeypatch, sse_events, **kwargs) -> list[str]:
+ """Drive a stream with a custom set of SSE events and return raw lines."""
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = _gemini_sse(sse_events),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ messages = kwargs.pop("messages", [{"role": "user", "content": "hi"}])
+ model = kwargs.pop("model", "gemini-2.5-flash")
+ temperature = kwargs.pop("temperature", 0.7)
+ top_p = kwargs.pop("top_p", 0.95)
+ max_tokens = kwargs.pop("max_tokens", 64)
+
+ out: list[str] = []
+
+ async def run():
+ client = _make_gemini_client()
+ async for line in client.stream_chat_completion(
+ messages = messages,
+ model = model,
+ temperature = temperature,
+ top_p = top_p,
+ max_tokens = max_tokens,
+ **kwargs,
+ ):
+ out.append(line)
+ await client.close()
+
+ _drive(run())
+ return out
+
+
+def _parse_chunks(lines: list[str]) -> list[dict]:
+ out: list[dict] = []
+ for raw in lines:
+ if not raw.startswith("data:"):
+ continue
+ payload = raw[len("data:") :].strip()
+ if not payload or payload == "[DONE]":
+ continue
+ try:
+ out.append(json.loads(payload))
+ except json.JSONDecodeError:
+ continue
+ return out
+
+
+# ── request body translation ─────────────────────────────────────────
+
+
+def test_request_body_uses_contents_and_parts_shape(monkeypatch):
+ """OpenAI messages must be translated to Gemini's `contents` shape."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "system", "content": "Be brief."},
+ {"role": "user", "content": "Hello"},
+ {"role": "assistant", "content": "Hi there"},
+ {"role": "user", "content": "Follow up"},
+ ],
+ )
+ body = captured["body"]
+ # system -> systemInstruction
+ assert body["systemInstruction"] == {"parts": [{"text": "Be brief."}]}, body
+ # user/assistant -> contents with role user/model
+ assert body["contents"] == [
+ {"role": "user", "parts": [{"text": "Hello"}]},
+ {"role": "model", "parts": [{"text": "Hi there"}]},
+ {"role": "user", "parts": [{"text": "Follow up"}]},
+ ], body["contents"]
+ # generationConfig fields map across with Google's casing.
+ gc = body["generationConfig"]
+ assert gc["temperature"] == 0.7
+ assert gc["topP"] == 0.95
+ assert gc["maxOutputTokens"] == 64
+
+
+def test_request_url_targets_stream_generate_content(monkeypatch):
+ """Helper must POST to /v1beta/models/{model}:streamGenerateContent?alt=sse."""
+ captured = _capture_body(monkeypatch, model = "gemini-2.5-pro")
+ url = captured["url"]
+ assert ":streamGenerateContent" in url, url
+ assert "alt=sse" in url, url
+ assert "/v1beta/models/gemini-2.5-pro" in url, url
+ assert captured["method"] == "POST"
+
+
+def test_request_auth_header_uses_x_goog_api_key(monkeypatch):
+ """API key must be sent on `x-goog-api-key`, not Authorization."""
+ captured = _capture_body(monkeypatch)
+ hdrs = captured["headers"]
+ assert hdrs.get("x-goog-api-key") == "AIza-test-key", hdrs
+ assert "authorization" not in {k.lower() for k in hdrs}, hdrs
+
+
+def test_top_k_forwarded_only_when_positive(monkeypatch):
+ """top_k is opt-in; only positive integers reach the wire."""
+ captured = _capture_body(monkeypatch, top_k = 40)
+ assert captured["body"]["generationConfig"]["topK"] == 40
+
+ captured = _capture_body(monkeypatch, top_k = 0)
+ assert "topK" not in captured["body"]["generationConfig"]
+
+
+def test_presence_penalty_forwarded_to_generation_config(monkeypatch):
+ """A non-zero presence_penalty reaches generationConfig.presencePenalty."""
+ captured = _capture_body(monkeypatch, presence_penalty = 0.7)
+ assert captured["body"]["generationConfig"]["presencePenalty"] == 0.7
+
+ # And the default of zero is omitted, matching top_k semantics.
+ captured = _capture_body(monkeypatch, presence_penalty = 0.0)
+ assert "presencePenalty" not in captured["body"]["generationConfig"]
+
+
+# ── thinkingConfig translation ────────────────────────────────────────
+
+
+def test_gemini25_flash_thinking_disabled_sets_budget_zero(monkeypatch):
+ """Gemini 2.5 Flash still uses thinkingBudget; 0 = off."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash",
+ enable_thinking = False,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingBudget": 0}, tc
+
+
+def test_gemini3_flash_thinking_disabled_uses_minimal_level(monkeypatch):
+ """Gemini 3 Flash migrated to thinkingLevel; "off" maps to minimal
+ (Gemini 3 cannot turn thinking fully off)."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3.5-flash",
+ enable_thinking = False,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": "minimal"}, tc
+
+
+def test_gemini25_pro_thinking_disabled_uses_small_budget(monkeypatch):
+ """Gemini 2.5 Pro 400s on thinkingBudget=0 ("only works in thinking
+ mode"); coerce to a small positive budget."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-pro",
+ enable_thinking = False,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc is not None and tc.get("thinkingBudget", 0) > 0, tc
+
+
+def test_gemini3_pro_thinking_disabled_uses_low_level(monkeypatch):
+ """Gemini 3 Pro uses thinkingLevel and rejects 'minimal' (Pro tier),
+ so 'off' coerces to 'low' (lowest the API accepts)."""
+ for model in (
+ "gemini-3.1-pro-preview",
+ "gemini-3-pro-preview",
+ "gemini-3.5-pro",
+ "gemini-pro-latest",
+ ):
+ captured = _capture_body(
+ monkeypatch,
+ model = model,
+ enable_thinking = False,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": "low"}, (model, tc)
+
+
+def test_gemini25_flash_effort_levels_map_to_budgets(monkeypatch):
+ """Gemini 2.5 Flash retains the integer thinkingBudget ladder."""
+ cases = {
+ "minimal": 512,
+ "low": 2048,
+ "medium": 8192,
+ "high": 24576,
+ "max": -1,
+ "xhigh": -1,
+ }
+ for effort, expected in cases.items():
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash",
+ reasoning_effort = effort,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingBudget": expected}, (effort, tc)
+
+
+def test_gemini3_flash_effort_levels_map_to_thinking_level(monkeypatch):
+ """Gemini 3 Flash thinkingLevel ladder: minimal/low/medium/high."""
+ cases = {
+ "minimal": "minimal",
+ "low": "low",
+ "medium": "medium",
+ "high": "high",
+ "max": "high",
+ }
+ for effort, expected in cases.items():
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3.5-flash",
+ reasoning_effort = effort,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": expected}, (effort, tc)
+
+
+def test_gemini3_pro_passes_medium_through(monkeypatch):
+ """Gemini 3.1+ Pro accepts thinkingLevel="medium" per
+ https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro;
+ forward as-is (medium is the documented mid-tier on Gemini 3.1)."""
+ for model in (
+ "gemini-3.1-pro-preview",
+ "gemini-pro-latest",
+ ):
+ captured = _capture_body(
+ monkeypatch,
+ model = model,
+ reasoning_effort = "medium",
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": "medium"}, (model, tc)
+
+
+def test_gemini3_pro_minimal_effort_coerces_to_low(monkeypatch):
+ """Gemini 3 Pro rejects thinkingLevel="minimal"; coerce to "low"."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3.1-pro-preview",
+ reasoning_effort = "minimal",
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": "low"}, tc
+
+
+def test_gemini3_flash_effort_none_maps_to_minimal(monkeypatch):
+ """reasoning_effort='none' on Gemini 3 Flash -> thinkingLevel=minimal."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3.5-flash",
+ reasoning_effort = "none",
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": "minimal"}, tc
+
+
+def test_thinking_default_omits_thinking_config(monkeypatch):
+ """When neither knob is supplied, thinkingConfig is omitted entirely
+ (Google's server-side default applies)."""
+ captured = _capture_body(monkeypatch, model = "gemini-3.5-flash")
+ gc = captured["body"]["generationConfig"]
+ assert "thinkingConfig" not in gc, gc
+
+
+def test_nano_banana_alias_routes_through_image_modalities(monkeypatch):
+ """`nano-banana-pro-preview` is an alias for the Pro image model and
+ must set responseModalities=[TEXT,IMAGE] when the Images pill is on
+ (enabled_tools includes "image_generation")."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "nano-banana-pro-preview",
+ enabled_tools = ["image_generation"],
+ )
+ gc = captured["body"]["generationConfig"]
+ assert gc.get("responseModalities") == ["TEXT", "IMAGE"], gc
+
+
+def test_image_capable_model_without_image_pill_stays_text_only(monkeypatch):
+ """When the Images pill is off (enabled_tools has no
+ image_generation), an image-capable model id (gemini-2.5-flash-image)
+ must force responseModalities=["TEXT"]. Google's image models
+ default to text+image when responseModalities is omitted, so
+ omitting it would silently bill image output the UI says is
+ disabled."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = [],
+ )
+ gc = captured["body"]["generationConfig"]
+ assert gc.get("responseModalities") == ["TEXT"], gc
+
+
+def test_image_models_skip_thinking_config(monkeypatch):
+ """Image-tier ids do not benefit from a visible thinking knob and
+ must NOT forward thinkingConfig even when stale UI state still
+ sends `reasoning_effort` or `enable_thinking=False`."""
+ for model in (
+ "gemini-2.5-flash-image",
+ "gemini-3.1-flash-image-preview",
+ "gemini-3-pro-image-preview",
+ "nano-banana-pro-preview",
+ ):
+ captured = _capture_body(
+ monkeypatch,
+ model = model,
+ reasoning_effort = "high",
+ enable_thinking = False,
+ enabled_tools = ["image_generation"],
+ )
+ gc = captured["body"]["generationConfig"]
+ assert "thinkingConfig" not in gc, (model, gc)
+
+
+def test_image_models_drop_code_execution(monkeypatch):
+ """All image-tier ids reject `tools: [{codeExecution: {}}]`; drop
+ silently. (Gemini 3 image models DO accept googleSearch -- see
+ test_gemini3_image_models_allow_google_search; older image models
+ drop everything.)"""
+ for model in (
+ "gemini-2.5-flash-image",
+ "gemini-3.1-flash-image-preview",
+ "gemini-3-pro-image-preview",
+ "nano-banana-pro-preview",
+ ):
+ captured = _capture_body(
+ monkeypatch,
+ model = model,
+ enabled_tools = ["image_generation", "code_execution"],
+ )
+ tools_arr = captured["body"].get("tools") or []
+ names = [list(t.keys())[0] for t in tools_arr]
+ assert "codeExecution" not in names, (model, tools_arr)
+
+
+def test_gemini_35_pro_uses_thinking_level(monkeypatch):
+ """`gemini-3.5-pro` is part of the Gemini 3 family and uses
+ thinkingLevel (not thinkingBudget). "Off" maps to "low" because Pro
+ tier rejects "minimal"."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3.5-pro",
+ enable_thinking = False,
+ )
+ tc = captured["body"]["generationConfig"].get("thinkingConfig")
+ assert tc == {"thinkingLevel": "low"}, tc
+
+
+def test_gemini3_image_models_allow_google_search(monkeypatch):
+ """Google documents Search grounding on the Gemini 3 image family
+ (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview,
+ nano-banana-pro). codeExecution stays blocked on image mode."""
+ for model in (
+ "gemini-3-pro-image-preview",
+ "gemini-3.1-flash-image-preview",
+ "nano-banana-pro-preview",
+ ):
+ captured = _capture_body(
+ monkeypatch,
+ model = model,
+ enabled_tools = ["image_generation", "web_search", "code_execution"],
+ )
+ tools_arr = captured["body"].get("tools") or []
+ names = [list(t.keys())[0] for t in tools_arr]
+ assert "googleSearch" in names, (model, tools_arr)
+ assert "codeExecution" not in names, (model, tools_arr)
+
+
+def test_legacy_image_models_block_google_search(monkeypatch):
+ """Older Gemini image ids (gemini-2.5-flash-image) still 400 on
+ `tools: [{googleSearch: {}}]`; backend keeps stripping it."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation", "web_search", "code_execution"],
+ )
+ assert "tools" not in captured["body"], captured["body"].get("tools")
+
+
+def test_legacy_openai_base_url_normalized(monkeypatch):
+ """Saved Gemini providers carrying the legacy `/v1beta/openai` base
+ (from the pre-PR OpenAI-compat plumbing) now point at the native
+ endpoint without the user re-saving the connection."""
+ client = ExternalProviderClient(
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta/openai",
+ api_key = "AIza-test-key",
+ )
+ assert client.base_url == "https://generativelanguage.googleapis.com/v1beta"
+
+
+def test_finish_reason_swaps_to_tool_calls_when_function_call_emitted(monkeypatch):
+ """Gemini emits finishReason="STOP" even for pure functionCall turns;
+ surface as `tool_calls` so OAI clients trigger tool execution."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {"functionCall": {"name": "lookup", "args": {"k": "v"}}}
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ]
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ finish_chunks = [
+ c for c in chunks if c.get("choices", [{}])[0].get("finish_reason") is not None
+ ]
+ assert finish_chunks, chunks
+ assert finish_chunks[-1]["choices"][0]["finish_reason"] == "tool_calls", chunks
+
+
+def test_thought_signature_round_trips_into_gemini_function_call(monkeypatch):
+ """An assistant tool_call carrying `extra_content.google.thought_signature`
+ must echo the value back as a sibling of the Gemini functionCall part."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "user", "content": "lookup x"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_0",
+ "type": "function",
+ "function": {"name": "lookup", "arguments": "{}"},
+ "extra_content": {"google": {"thought_signature": "SIG-ABC"}},
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_0",
+ "name": "lookup",
+ "content": "{}",
+ },
+ ],
+ )
+ contents = captured["body"]["contents"]
+ fc_turn = next((c for c in contents if c["role"] == "model"), None)
+ assert fc_turn is not None, contents
+ fc_part = next(
+ (p for p in fc_turn["parts"] if "functionCall" in p),
+ None,
+ )
+ assert fc_part is not None, fc_turn
+ assert fc_part.get("thoughtSignature") == "SIG-ABC", fc_part
+
+
+def test_thought_signature_emitted_in_tool_call_delta(monkeypatch):
+ """A Gemini functionCall part with `thoughtSignature` must surface
+ that signature on the outbound OpenAI tool_calls delta via
+ `extra_content.google.thought_signature`."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "functionCall": {
+ "name": "lookup",
+ "args": {"k": "v"},
+ "id": "call_xyz",
+ },
+ "thoughtSignature": "SIG-FROM-GEMINI",
+ }
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ]
+ }
+ ]
+ chunks = _parse_chunks(_collect(monkeypatch, sse))
+ deltas = [
+ tc
+ for c in chunks
+ for tc in (c.get("choices", [{}])[0].get("delta", {}) or {}).get(
+ "tool_calls", []
+ )
+ ]
+ assert deltas, chunks
+ sig = deltas[0].get("extra_content", {}).get("google", {}).get("thought_signature")
+ assert sig == "SIG-FROM-GEMINI", deltas
+
+
+def test_image_models_suppress_phantom_web_search_card(monkeypatch):
+ """When the image guard filters googleSearch out of the outbound
+ request, the inbound stream must NOT emit web_search tool_start /
+ tool_end (otherwise the UI shows a misleading 'Search complete'
+ card on a turn where Gemini never actually searched)."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {"role": "model", "parts": [{"text": "drawn"}]},
+ "finishReason": "STOP",
+ }
+ ]
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation", "web_search", "code_execution"],
+ )
+ chunks = _parse_chunks(lines)
+ tool_evs = [
+ ev
+ for c in chunks
+ for ev in [c.get("_toolEvent")]
+ if isinstance(ev, dict) and ev.get("tool_name") == "web_search"
+ ]
+ assert tool_evs == [], tool_evs
+
+
+def test_image_generation_tool_on_image_model_drops_text_tools(monkeypatch):
+ """`enabled_tools=["image_generation", "web_search", "code_execution"]`
+ on a Gemini IMAGE model flips responseModalities to TEXT+IMAGE; in
+ that mode codeExecution must NOT be forwarded (Gemini rejects text
+ code tools alongside image responseModalities). Older image
+ families also drop googleSearch."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = [
+ "image_generation",
+ "web_search",
+ "code_execution",
+ ],
+ )
+ assert "tools" not in captured["body"], captured["body"]
+ assert captured["body"]["generationConfig"].get("responseModalities") == [
+ "TEXT",
+ "IMAGE",
+ ]
+
+
+def test_prompt_feedback_block_reason_surfaces_as_error(monkeypatch):
+ """`promptFeedback.blockReason` with zero candidates must produce
+ an error chunk, not a silent empty assistant reply."""
+ sse = [
+ {
+ "promptFeedback": {"blockReason": "SAFETY"},
+ }
+ ]
+ chunks = _parse_chunks(_collect(monkeypatch, sse))
+ error_chunks = [c for c in chunks if "error" in c]
+ assert error_chunks, chunks
+ assert "SAFETY" in (
+ error_chunks[0].get("error", {}).get("message") or ""
+ ), error_chunks
+
+
+def test_usage_chunk_includes_thoughts_tokens(monkeypatch):
+ """`thoughtsTokenCount` is the hidden-reasoning slice of output;
+ roll it into `output_tokens` AND surface it on
+ `output_tokens_details.reasoning_tokens` so total_tokens reflects
+ the full billable spend."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {"role": "model", "parts": [{"text": "ok"}]},
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 10,
+ "candidatesTokenCount": 5,
+ "thoughtsTokenCount": 20,
+ "totalTokenCount": 35,
+ },
+ }
+ ]
+ chunks = _parse_chunks(_collect(monkeypatch, sse))
+ usage_chunk = next((c for c in chunks if isinstance(c.get("usage"), dict)), None)
+ assert usage_chunk is not None, chunks
+ usage = usage_chunk["usage"]
+ assert usage.get("prompt_tokens") == 10, usage
+ # candidates 5 + thoughts 20 = 25 output tokens; total = 35.
+ assert usage.get("completion_tokens") == 25, usage
+ assert usage.get("total_tokens") == 35, usage
+
+
+# ── web_search forwarded as googleSearch tool ────────────────────────
+
+
+def test_web_search_forwarded_as_google_search_tool(monkeypatch):
+ captured = _capture_body(
+ monkeypatch,
+ enabled_tools = ["web_search"],
+ )
+ tools = captured["body"].get("tools") or []
+ assert {"googleSearch": {}} in tools, tools
+
+
+def test_code_execution_forwarded_as_code_execution_tool(monkeypatch):
+ captured = _capture_body(
+ monkeypatch,
+ enabled_tools = ["code_execution"],
+ )
+ tools = captured["body"].get("tools") or []
+ assert {"codeExecution": {}} in tools, tools
+
+
+def test_omitted_tools_leaves_body_untouched(monkeypatch):
+ captured = _capture_body(monkeypatch, enabled_tools = [])
+ assert "tools" not in captured["body"], captured["body"]
+
+
+# ── prompt caching passthrough ───────────────────────────────────────
+
+
+def test_cached_content_pass_through(monkeypatch):
+ """A string cache id on enable_prompt_caching is forwarded verbatim."""
+ cache_name = "cachedContents/abc123"
+ captured = _capture_body(
+ monkeypatch,
+ enable_prompt_caching = cache_name,
+ )
+ assert captured["body"].get("cachedContent") == cache_name
+
+
+def test_boolean_caching_does_not_set_cached_content(monkeypatch):
+ """Studio's existing True/False signals shouldn't fabricate a cache id."""
+ captured = _capture_body(monkeypatch, enable_prompt_caching = True)
+ assert "cachedContent" not in captured["body"]
+
+
+# ── image generation: request modalities + response translation ──────
+
+
+def test_image_model_sets_response_modalities(monkeypatch):
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation"],
+ )
+ assert captured["body"]["generationConfig"]["responseModalities"] == [
+ "TEXT",
+ "IMAGE",
+ ]
+
+
+def test_image_generation_tool_sets_response_modalities_on_image_model(monkeypatch):
+ """`enabled_tools=["image_generation"]` flips responseModalities
+ only when the selected model is image-capable; otherwise the
+ request stays plain text (text-only models 400 on
+ responseModalities)."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation"],
+ )
+ assert captured["body"]["generationConfig"]["responseModalities"] == [
+ "TEXT",
+ "IMAGE",
+ ]
+
+
+def test_image_response_emits_image_b64_tool_event(monkeypatch):
+ """`inlineData` parts become a tool_end with image_b64 + image_mime."""
+ fake_b64 = base64.b64encode(b"PNG-BYTES").decode()
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": fake_b64,
+ }
+ }
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 5,
+ "candidatesTokenCount": 0,
+ },
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ model = "gemini-2.5-flash-image",
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ starts = [e for e in tool_events if e.get("type") == "tool_start"]
+ ends = [e for e in tool_events if e.get("type") == "tool_end"]
+ image_starts = [e for e in starts if e.get("tool_name") == "image_generation"]
+ image_ends = [e for e in ends if e.get("image_b64")]
+ assert len(image_starts) == 1, tool_events
+ assert len(image_ends) == 1, tool_events
+ assert image_ends[0]["image_b64"] == fake_b64
+ assert image_ends[0]["image_mime"] == "image/png"
+
+
+# ── function calling round-trips both directions ─────────────────────
+
+
+def test_function_call_response_translates_to_tool_calls_delta(monkeypatch):
+ """Gemini `functionCall` parts become OpenAI `tool_calls` delta chunks."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "functionCall": {
+ "name": "get_weather",
+ "args": {"location": "Paris"},
+ }
+ }
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 12,
+ "candidatesTokenCount": 4,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ tool_call_chunks = [
+ c
+ for c in chunks
+ if "_toolEvent" not in c
+ and any(
+ (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"])
+ for ch in c.get("choices", [])
+ )
+ ]
+ assert len(tool_call_chunks) == 1, chunks
+ tc = tool_call_chunks[0]["choices"][0]["delta"]["tool_calls"][0]
+ assert tc["function"]["name"] == "get_weather"
+ args = json.loads(tc["function"]["arguments"])
+ assert args == {"location": "Paris"}
+
+
+def test_tool_message_translates_to_function_response_part(monkeypatch):
+ """role=tool follow-ups are rewritten to functionResponse parts."""
+ messages = [
+ {"role": "user", "content": "Weather?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_1",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": json.dumps({"location": "Paris"}),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "name": "get_weather",
+ "content": json.dumps({"temp_c": 18, "summary": "Sunny"}),
+ },
+ ]
+ captured = _capture_body(monkeypatch, messages = messages)
+ contents = captured["body"]["contents"]
+ # Last turn must be a functionResponse part (Gemini wraps it as a
+ # role=user turn carrying the result).
+ last = contents[-1]
+ assert last["role"] == "user", last
+ fr = last["parts"][0].get("functionResponse")
+ assert fr is not None, last
+ assert fr["name"] == "get_weather"
+ assert fr["response"] == {"temp_c": 18, "summary": "Sunny"}
+ # And the assistant turn carries the original functionCall so the
+ # model sees the round-trip context.
+ assistant_turn = [c for c in contents if c["role"] == "model"][0]
+ fc_part = next(
+ (p for p in assistant_turn["parts"] if "functionCall" in p),
+ None,
+ )
+ assert fc_part is not None, assistant_turn
+ assert fc_part["functionCall"]["name"] == "get_weather"
+ assert fc_part["functionCall"]["args"] == {"location": "Paris"}
+
+
+def test_parallel_function_calls_get_distinct_tool_call_indices(monkeypatch):
+ """Each emitted functionCall in one assistant turn needs its own
+ tool_calls[*].index. Hardcoding index=0 collapses parallel calls
+ onto a single slot in OpenAI-style reassemblers."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "functionCall": {
+ "id": "call_alpha",
+ "name": "search",
+ "args": {"q": "alpha"},
+ }
+ },
+ {
+ "functionCall": {
+ "id": "call_beta",
+ "name": "search",
+ "args": {"q": "beta"},
+ }
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 8,
+ "candidatesTokenCount": 4,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ tool_call_chunks = [
+ c
+ for c in chunks
+ if "_toolEvent" not in c
+ and any(
+ (isinstance(ch.get("delta"), dict) and "tool_calls" in ch["delta"])
+ for ch in c.get("choices", [])
+ )
+ ]
+ assert len(tool_call_chunks) == 2, tool_call_chunks
+ indices = [
+ c["choices"][0]["delta"]["tool_calls"][0]["index"] for c in tool_call_chunks
+ ]
+ assert indices == [0, 1], indices
+
+
+def test_function_call_ids_forwarded_into_gemini_function_call_part(monkeypatch):
+ """OpenAI tool_call id rides functionCall.id so parallel calls disambiguate."""
+ messages = [
+ {"role": "user", "content": "x"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_alpha",
+ "type": "function",
+ "function": {
+ "name": "search",
+ "arguments": json.dumps({"q": "a"}),
+ },
+ },
+ {
+ "id": "call_beta",
+ "type": "function",
+ "function": {
+ "name": "search",
+ "arguments": json.dumps({"q": "b"}),
+ },
+ },
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_alpha",
+ "content": json.dumps({"hits": ["A"]}),
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_beta",
+ "content": json.dumps({"hits": ["B"]}),
+ },
+ ]
+ captured = _capture_body(monkeypatch, messages = messages)
+ contents = captured["body"]["contents"]
+ assistant_parts = next(c for c in contents if c["role"] == "model")["parts"]
+ call_ids = [p["functionCall"]["id"] for p in assistant_parts if "functionCall" in p]
+ assert call_ids == ["call_alpha", "call_beta"], assistant_parts
+ response_ids = [
+ p["functionResponse"]["id"]
+ for c in contents
+ for p in c["parts"]
+ if "functionResponse" in p
+ ]
+ assert response_ids == ["call_alpha", "call_beta"], contents
+
+
+def test_parse_gemini_models_translates_native_catalog():
+ """Gemini's native /v1beta/models payload becomes OpenAI-shape entries."""
+ payload = {
+ "models": [
+ {
+ "name": "models/gemini-2.5-flash",
+ "baseModelId": "gemini-2.5-flash",
+ "displayName": "Gemini 2.5 Flash",
+ "supportedGenerationMethods": [
+ "generateContent",
+ "streamGenerateContent",
+ ],
+ },
+ {
+ "name": "models/embedding-001",
+ "supportedGenerationMethods": ["embedContent"],
+ },
+ {
+ "name": "models/gemini-2.5-pro",
+ },
+ ]
+ }
+ out = ExternalProviderClient._parse_gemini_models(payload)
+ ids = [m["id"] for m in out]
+ assert "gemini-2.5-flash" in ids
+ assert "gemini-2.5-pro" in ids
+ assert "embedding-001" not in ids
+ flash = next(m for m in out if m["id"] == "gemini-2.5-flash")
+ assert flash["display_name"] == "Gemini 2.5 Flash"
+ assert flash["owned_by"] == "google"
+
+
+def test_code_execution_parts_translate_to_code_execution_tool_events(monkeypatch):
+ """executableCode + codeExecutionResult parts emit code_execution events."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "executableCode": {
+ "language": "PYTHON",
+ "code": "print(2+2)",
+ }
+ },
+ {
+ "codeExecutionResult": {
+ "outcome": "OUTCOME_OK",
+ "output": "4\n",
+ }
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 8,
+ "candidatesTokenCount": 4,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"])
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ code_starts = [
+ e
+ for e in tool_events
+ if e.get("type") == "tool_start" and e.get("tool_name") == "code_execution"
+ ]
+ code_ends = [
+ e
+ for e in tool_events
+ if e.get("type") == "tool_end" and "4" in str(e.get("result", ""))
+ ]
+ assert len(code_starts) == 1, tool_events
+ assert code_starts[0]["arguments"]["code"] == "print(2+2)"
+ assert code_starts[0]["arguments"]["language"] == "python"
+ assert len(code_ends) == 1, tool_events
+ # tool_start and tool_end must share the same tool_call_id so the
+ # frontend pairs them onto a single CodeExecutionToolUI block.
+ assert code_starts[0]["tool_call_id"] == code_ends[0]["tool_call_id"]
+
+
+def test_code_execution_failure_outcome_surfaces_in_result(monkeypatch):
+ """OUTCOME_FAILED is prefixed onto the result text so the UI shows it."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "executableCode": {
+ "language": "PYTHON",
+ "code": "1/0",
+ }
+ },
+ {
+ "codeExecutionResult": {
+ "outcome": "OUTCOME_FAILED",
+ "output": "ZeroDivisionError",
+ }
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 5,
+ "candidatesTokenCount": 2,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse, enabled_tools = ["code_execution"])
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ result_text = next(
+ (e["result"] for e in tool_events if e.get("type") == "tool_end"),
+ "",
+ )
+ assert "OUTCOME_FAILED" in result_text
+ assert "ZeroDivisionError" in result_text
+
+
+def test_tool_message_recovers_name_from_tool_call_id(monkeypatch):
+ """When name is omitted, recover it from the matching tool_call_id."""
+ messages = [
+ {"role": "user", "content": "Weather?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_xyz",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": json.dumps({"location": "Paris"}),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_xyz",
+ "content": json.dumps({"temp_c": 18}),
+ },
+ ]
+ captured = _capture_body(monkeypatch, messages = messages)
+ contents = captured["body"]["contents"]
+ last = contents[-1]
+ fr = last["parts"][0].get("functionResponse")
+ assert fr is not None, last
+ assert (
+ fr["name"] == "get_weather"
+ ), "name should fall back to the prior tool_call's function name"
+
+
+# ── usage chunk surfaces promptTokenCount / candidatesTokenCount ─────
+
+
+def test_usage_chunk_translates_gemini_token_counts(monkeypatch):
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "ok"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 1234,
+ "candidatesTokenCount": 56,
+ "cachedContentTokenCount": 1000,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ usage_chunks = [c for c in chunks if c.get("choices") == [] and "usage" in c]
+ assert len(usage_chunks) == 1, chunks
+ usage = usage_chunks[0]["usage"]
+ assert usage["prompt_tokens"] == 1234
+ assert usage["completion_tokens"] == 56
+ assert usage["total_tokens"] == 1290
+ assert usage["prompt_tokens_details"]["cached_tokens"] == 1000
+
+
+# ── multimodal: vision image -> inlineData ───────────────────────────
+
+
+def test_vision_data_url_translates_to_inline_data(monkeypatch):
+ fake = base64.b64encode(b"JPGBYTES").decode()
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What is this?"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": f"data:image/jpeg;base64,{fake}",
+ },
+ },
+ ],
+ }
+ ]
+ captured = _capture_body(monkeypatch, messages = messages)
+ parts = captured["body"]["contents"][0]["parts"]
+ inline_parts = [p for p in parts if "inlineData" in p]
+ assert len(inline_parts) == 1, parts
+ assert inline_parts[0]["inlineData"] == {
+ "mimeType": "image/jpeg",
+ "data": fake,
+ }
+
+
+# ── finish reason mapping ────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+ "gemini_reason, openai_reason",
+ [
+ ("STOP", "stop"),
+ ("MAX_TOKENS", "length"),
+ ("SAFETY", "content_filter"),
+ ("PROHIBITED_CONTENT", "content_filter"),
+ ],
+)
+def test_finish_reason_translation(monkeypatch, gemini_reason, openai_reason):
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "x"}],
+ },
+ "finishReason": gemini_reason,
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 1,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ finish_chunks = [
+ c for c in chunks if any(ch.get("finish_reason") for ch in c.get("choices", []))
+ ]
+ assert any(
+ ch["choices"][0]["finish_reason"] == openai_reason for ch in finish_chunks
+ ), finish_chunks
+
+
+# ── grounding citations surface as web_search tool_end ───────────────
+
+
+def test_grounding_metadata_surfaces_as_tool_end_citations(monkeypatch):
+ """`groundingMetadata.groundingChunks[].web` -> tool_end result block."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "Answer with sources."}],
+ },
+ "groundingMetadata": {
+ "groundingChunks": [
+ {
+ "web": {
+ "uri": "https://example.com/a",
+ "title": "Example A",
+ }
+ },
+ {
+ "web": {
+ "uri": "https://example.com/b",
+ "title": "Example B",
+ }
+ },
+ ]
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 7,
+ "candidatesTokenCount": 3,
+ },
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ enabled_tools = ["web_search"],
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ web_search_ends = [
+ e
+ for e in tool_events
+ if e.get("type") == "tool_end" and e.get("tool_call_id") == "gemini_web_search"
+ ]
+ assert len(web_search_ends) == 1, tool_events
+ result = web_search_ends[0]["result"]
+ assert "https://example.com/a" in result
+ assert "https://example.com/b" in result
+ assert "Example A" in result
+ assert "Example B" in result
+
+
+# ── round 3 review follow-ups ─────────────────────────────────────────
+
+
+def test_custom_gemini_proxy_base_url_not_rewritten():
+ """Only the Google-hosted /v1beta/openai base is normalized; a
+ custom gateway whose path ends in /openai must be left alone."""
+ client = ExternalProviderClient(
+ provider_type = "gemini",
+ base_url = "https://proxy.example.com/team/openai",
+ api_key = "AIza-test-key",
+ )
+ assert client.base_url == "https://proxy.example.com/team/openai"
+
+
+def test_custom_gemini_proxy_uses_openai_dispatch():
+ """Any non-Google Gemini base (LiteLLM, custom OpenAI-compat
+ routers) must route through the OpenAI-compatible forwarder, not
+ the native translator. Auth uses Authorization: Bearer ..., not
+ x-goog-api-key."""
+ for base in (
+ "https://proxy.example.com/team/openai",
+ "https://proxy.example.com/v1",
+ "https://litellm.internal.example/v1",
+ ):
+ client = ExternalProviderClient(
+ provider_type = "gemini",
+ base_url = base,
+ api_key = "AIza-test-key",
+ )
+ assert client._is_openai_compatible() is True, base
+ headers = client._auth_headers()
+ assert "x-goog-api-key" not in {k.lower() for k in headers}, (
+ base,
+ headers,
+ )
+ assert headers["Authorization"] == "Bearer AIza-test-key", (
+ base,
+ headers,
+ )
+
+
+def test_google_hosted_gemini_still_uses_native_dispatch():
+ """Google-hosted Gemini keeps native dispatch + x-goog-api-key auth."""
+ client = ExternalProviderClient(
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta",
+ api_key = "AIza-test-key",
+ )
+ assert client._is_openai_compatible() is False
+ headers = client._auth_headers()
+ assert headers.get("x-goog-api-key") == "AIza-test-key", headers
+
+
+def test_invalid_gemini_model_id_rejected_before_request(monkeypatch):
+ """Path-traversal model ids must be rejected before the URL is
+ interpolated so the configured API key isn't sent to unintended
+ Gemini endpoints."""
+
+ captured: list[httpx.Request] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured.append(request)
+ return httpx.Response(
+ 200,
+ content = _gemini_sse([]),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ out: list[str] = []
+
+ async def run():
+ client = _make_gemini_client()
+ async for line in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "../cachedContents/leak",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ ):
+ out.append(line)
+ await client.close()
+
+ _drive(run())
+ # No outbound request should have been issued.
+ assert captured == [], captured
+ error_lines = [line for line in out if '"error"' in line]
+ assert error_lines, out
+
+
+def test_top_k_omitted_when_not_explicit_default_for_gemini(monkeypatch):
+ """top_k=None means "use provider default"; helper must not emit
+ `topK` in generationConfig when the caller didn't pass it."""
+ captured = _capture_body(monkeypatch, top_k = None)
+ assert "topK" not in captured["body"]["generationConfig"], captured["body"]
+
+
+def test_text_model_image_generation_tool_silently_dropped(monkeypatch):
+ """A stale `enabled_tools=["image_generation"]` on a text-only
+ Gemini model (e.g. gemini-2.5-flash) must NOT switch the request
+ into image mode -- Google's API 400s on responseModalities for
+ text models."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash",
+ enabled_tools = ["image_generation"],
+ )
+ gc = captured["body"]["generationConfig"]
+ assert "responseModalities" not in gc, gc
+
+
+def test_empty_text_part_with_thought_signature_emits_extra_content(
+ monkeypatch,
+):
+ """Gemini 3 can ship a content-free fragment whose only payload is
+ `thoughtSignature`. The translator must still surface that signature
+ on a delta.extra_content envelope so the next turn can replay it."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {"text": "answer"},
+ {"thoughtSignature": "SIG-FINAL"},
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 2,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ extra_carriers = [
+ c
+ for c in chunks
+ if c.get("choices")
+ and c["choices"][0]["delta"].get("extra_content")
+ == {"google": {"thought_signature": "SIG-FINAL"}}
+ ]
+ assert extra_carriers, chunks
+
+
+def test_enable_prompt_caching_false_string_coerces_to_bool():
+ """Pre-PR the field was Optional[bool]; widening to Union[bool,str]
+ must preserve historical coercion so callers sending `"false"`
+ still opt out of caching."""
+ from models.inference import ChatCompletionRequest
+
+ msg = {"role": "user", "content": "hi"}
+ req = ChatCompletionRequest.model_validate(
+ {
+ "model": "gemini-2.5-flash",
+ "messages": [msg],
+ "enable_prompt_caching": "false",
+ }
+ )
+ assert req.enable_prompt_caching is False, req.enable_prompt_caching
+
+ req = ChatCompletionRequest.model_validate(
+ {
+ "model": "gemini-2.5-flash",
+ "messages": [msg],
+ "enable_prompt_caching": "true",
+ }
+ )
+ assert req.enable_prompt_caching is True
+
+ # An actual cache resource name passes through untouched.
+ req = ChatCompletionRequest.model_validate(
+ {
+ "model": "gemini-2.5-flash",
+ "messages": [msg],
+ "enable_prompt_caching": "cachedContents/abc123",
+ }
+ )
+ assert req.enable_prompt_caching == "cachedContents/abc123"
+
+
+def test_legacy_google_openai_base_url_is_rewritten():
+ """The Google-hosted /v1beta/openai legacy base IS still rewritten."""
+ client = ExternalProviderClient(
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta/openai",
+ api_key = "AIza-test-key",
+ )
+ assert client.base_url == "https://generativelanguage.googleapis.com/v1beta"
+
+
+def test_remote_image_url_downloads_and_inlines_as_base64(monkeypatch):
+ """Round 14: arbitrary public HTTPS image URLs cannot be sent as
+ Gemini fileData (that path is reserved for Files API URIs and
+ YouTube). The translator must fetch the bytes server-side and
+ inline them as base64 inlineData."""
+ image_bytes = b"FAKEPNGBYTES"
+
+ async def fake_fetch(url, fallback_mime, max_bytes = None):
+ assert url == "https://cdn.example.com/diagram.png"
+ return ("image/png", base64.b64encode(image_bytes).decode("ascii"))
+
+ monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch)
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "what is this?"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://cdn.example.com/diagram.png",
+ },
+ },
+ ],
+ }
+ ],
+ )
+ parts = captured["body"]["contents"][-1]["parts"]
+ inline = next((p for p in parts if "inlineData" in p), None)
+ assert inline is not None, parts
+ assert inline["inlineData"]["mimeType"] == "image/png"
+ assert inline["inlineData"]["data"] == base64.b64encode(image_bytes).decode()
+ assert not any("fileData" in p for p in parts), parts
+
+
+def test_remote_image_url_dropped_when_fetch_returns_none(monkeypatch):
+ """Round 15: if the SSRF guard rejects the URL (private host,
+ non-https, oversize, non-image), the helper returns None and the
+ image part is silently dropped instead of forwarding raw bytes
+ or a fileData fallback."""
+
+ async def fake_fetch_reject(url, fallback_mime, max_bytes = None):
+ return None
+
+ monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch_reject)
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "what is this?"},
+ {
+ "type": "image_url",
+ "image_url": {"url": "http://10.0.0.5/private.png"},
+ },
+ ],
+ }
+ ],
+ )
+ parts = captured["body"]["contents"][-1]["parts"]
+ assert not any("inlineData" in p for p in parts), parts
+ assert not any("fileData" in p for p in parts), parts
+
+
+def test_safe_fetch_image_rejects_non_https():
+ """SSRF guard: only https URLs may be fetched."""
+ res = asyncio.new_event_loop().run_until_complete(
+ ep_mod._safe_fetch_image_for_gemini("http://cdn.example.com/x.png", "image/png")
+ )
+ assert res is None
+
+
+def test_safe_fetch_image_rejects_loopback_ip_literal():
+ """SSRF guard: refuse loopback / private IP literals before any
+ network call."""
+ for url in (
+ "https://127.0.0.1/x.png",
+ "https://[::1]/x.png",
+ "https://169.254.169.254/latest/meta-data",
+ "https://10.0.0.5/x.png",
+ "https://192.168.1.1/x.png",
+ ):
+ res = asyncio.new_event_loop().run_until_complete(
+ ep_mod._safe_fetch_image_for_gemini(url, "image/png")
+ )
+ assert res is None, url
+
+
+def test_safe_fetch_image_rejects_resolved_private_host(monkeypatch):
+ """SSRF guard: if a hostname resolves to a private IP, refuse."""
+ import socket
+
+ def fake_getaddrinfo(host, *_args, **_kwargs):
+ return [(socket.AF_INET, None, None, "", ("10.0.0.5", 0))]
+
+ monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
+ res = asyncio.new_event_loop().run_until_complete(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://internal.example/x.png", "image/png"
+ )
+ )
+ assert res is None
+
+
+def test_youtube_and_files_api_uris_stay_as_file_data(monkeypatch):
+ """Round 14: YouTube URLs and generativelanguage.googleapis.com
+ Files API URIs are the documented `fileData.fileUri` paths and
+ must NOT be downloaded; arbitrary public URLs do get fetched."""
+ captured: dict = {}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = _gemini_sse(
+ [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "ok"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 1,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ ),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = _make_gemini_client()
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "explain"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://www.youtube.com/watch?v=abc123",
+ },
+ },
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://generativelanguage.googleapis.com/v1beta/files/abc",
+ },
+ },
+ ],
+ }
+ ],
+ model = "gemini-2.5-flash",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 64,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ parts = captured["body"]["contents"][-1]["parts"]
+ file_uris = [p["fileData"]["fileUri"] for p in parts if "fileData" in p]
+ assert "https://www.youtube.com/watch?v=abc123" in file_uris, parts
+ assert (
+ "https://generativelanguage.googleapis.com/v1beta/files/abc" in file_uris
+ ), parts
+
+
+def test_tool_use_prompt_tokens_added_to_input_tokens(monkeypatch):
+ """`toolUsePromptTokenCount` must roll into the OpenAI prompt
+ total -- otherwise tool turns silently undercount input tokens."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "result"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 10,
+ "toolUsePromptTokenCount": 100,
+ "candidatesTokenCount": 5,
+ "thoughtsTokenCount": 2,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ usage_chunks = [c for c in chunks if c.get("usage")]
+ assert len(usage_chunks) == 1, chunks
+ usage = usage_chunks[0]["usage"]
+ assert usage["prompt_tokens"] == 110, usage
+ assert usage["completion_tokens"] == 7, usage
+ assert usage["total_tokens"] == 117, usage
+ assert usage["completion_tokens_details"]["reasoning_tokens"] == 2, usage
+
+
+def test_usage_chunk_reasoning_tokens_surfaced(monkeypatch):
+ """thoughtsTokenCount must surface as completion_tokens_details.
+ reasoning_tokens in the emitted OpenAI usage chunk."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "ok"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 8,
+ "candidatesTokenCount": 5,
+ "thoughtsTokenCount": 20,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ usage_chunks = [c for c in chunks if c.get("usage")]
+ assert len(usage_chunks) == 1, chunks
+ usage = usage_chunks[0]["usage"]
+ assert usage["completion_tokens"] == 25, usage
+ assert usage["completion_tokens_details"]["reasoning_tokens"] == 20, usage
+
+
+def test_prompt_block_pairs_web_search_tool_end(monkeypatch):
+ """When `promptFeedback.blockReason` triggers after the synthetic
+ web_search tool_start, the helper must emit a matching tool_end so
+ the UI does not leave a "searching..." spinner stuck on screen."""
+ sse = [
+ {"promptFeedback": {"blockReason": "SAFETY"}},
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ enabled_tools = ["web_search"],
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ starts = [e for e in tool_events if e.get("type") == "tool_start"]
+ ends = [e for e in tool_events if e.get("type") == "tool_end"]
+ assert len(starts) == 1, tool_events
+ assert len(ends) == 1, tool_events
+ assert ends[0]["tool_call_id"] == "gemini_web_search"
+ assert "aborted" in ends[0]["result"]
+ error_chunks = [c for c in chunks if c.get("error")]
+ assert error_chunks, chunks
+
+
+def test_code_execution_tool_events_stow_native_part(monkeypatch):
+ """executableCode / codeExecutionResult must round-trip native ids
+ and thoughtSignature in google.native_part so follow-up turns can
+ replay Gemini's required history shape."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "executableCode": {
+ "id": "code_a",
+ "language": "PYTHON",
+ "code": "print(1+1)",
+ },
+ "thoughtSignature": "SIG-CODE",
+ },
+ {
+ "codeExecutionResult": {
+ "id": "result_a",
+ "outcome": "OUTCOME_OK",
+ "output": "2\n",
+ },
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 5,
+ "candidatesTokenCount": 4,
+ },
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ enabled_tools = ["code_execution"],
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ starts = [e for e in tool_events if e.get("type") == "tool_start"]
+ ends = [e for e in tool_events if e.get("type") == "tool_end"]
+ code_start = next(
+ (e for e in starts if e.get("tool_name") == "code_execution"),
+ None,
+ )
+ code_end = next(iter(ends), None)
+ assert code_start is not None, starts
+ assert code_start["tool_call_id"] == "code_a", code_start
+ native = code_start["arguments"]["google"]["native_part"]
+ # Round 21: native_part now uses an ordered `parts` list so per-part
+ # `thoughtSignature` survives a frontend merge of executableCode +
+ # codeExecutionResult into one tool-call card.
+ start_parts = native["parts"]
+ assert start_parts[0]["executableCode"]["id"] == "code_a"
+ assert start_parts[0]["thoughtSignature"] == "SIG-CODE"
+ assert code_end is not None, ends
+ assert code_end["tool_call_id"] == "code_a", code_end
+ native_end = code_end["google"]["native_part"]
+ end_parts = native_end["parts"]
+ assert end_parts[0]["codeExecutionResult"]["id"] == "result_a"
+
+
+def test_inline_image_tool_end_carries_thought_signature(monkeypatch):
+ """Inline image parts with thoughtSignature must persist it on the
+ emitted tool_end so Gemini 3 image editing can echo it back."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": base64.b64encode(b"PNG").decode(),
+ },
+ "thoughtSignature": "SIG-IMG",
+ }
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 4,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ model = "gemini-2.5-flash-image",
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ image_ends = [
+ e for e in tool_events if e.get("type") == "tool_end" and e.get("image_b64")
+ ]
+ assert image_ends, tool_events
+ assert image_ends[0]["google"]["thought_signature"] == "SIG-IMG"
+ # Multi-turn image edit must replay the original inlineData part with
+ # its thoughtSignature; the outbound translator reads
+ # google.native_part.parts[].inlineData, so stow it on the tool_end
+ # too. Round 21 changed native_part to an ordered parts list so a
+ # per-part signature stays attached to inlineData only.
+ native = image_ends[0]["google"]["native_part"]
+ image_parts = native["parts"]
+ assert image_parts[0]["inlineData"]["mimeType"] == "image/png"
+ assert image_parts[0]["inlineData"]["data"] == base64.b64encode(b"PNG").decode()
+ assert image_parts[0]["thoughtSignature"] == "SIG-IMG"
+
+
+def test_code_execution_plot_attaches_inline_image_native_part(monkeypatch):
+ """A code_execution turn that returns a matplotlib plot must stow
+ the plot's inlineData on the secondary tool_end so the follow-up
+ turn can replay the image alongside executableCode and
+ codeExecutionResult."""
+ plot_data = base64.b64encode(b"PLOT").decode()
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "executableCode": {
+ "id": "code_a",
+ "language": "PYTHON",
+ "code": "plt.plot([0,1])",
+ },
+ },
+ {
+ "codeExecutionResult": {
+ "id": "result_a",
+ "outcome": "OUTCOME_OK",
+ "output": "",
+ },
+ },
+ {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": plot_data,
+ },
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 5,
+ "candidatesTokenCount": 4,
+ },
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ enabled_tools = ["code_execution"],
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ code_ends = [
+ e
+ for e in tool_events
+ if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_a"
+ ]
+ # Two tool_end events on the same id: one for codeExecutionResult,
+ # one merging in the inlineData plot. The plot one must carry the
+ # native inlineData under google.native_part so the frontend
+ # tool_end merge union joins it with the prior executableCode and
+ # codeExecutionResult parts on the same card.
+ assert len(code_ends) == 2, code_ends
+ image_end = next(
+ (e for e in code_ends if "__IMAGES__:" in (e.get("result") or "")),
+ None,
+ )
+ assert image_end is not None, code_ends
+ native = image_end["google"]["native_part"]
+ plot_parts = native["parts"]
+ assert plot_parts[0]["inlineData"]["mimeType"] == "image/png"
+ assert plot_parts[0]["inlineData"]["data"] == plot_data
+
+
+def test_text_chunk_carries_thought_signature(monkeypatch):
+ """Text parts with thoughtSignature surface it on delta.extra_content
+ so frontend persistence can replay it on the follow-up turn."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "text": "hello",
+ "thoughtSignature": "SIG-TEXT",
+ }
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 2,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ lines = _collect(monkeypatch, sse)
+ chunks = _parse_chunks(lines)
+ text_chunks = [
+ c
+ for c in chunks
+ if c.get("choices") and c["choices"][0]["delta"].get("content") == "hello"
+ ]
+ assert text_chunks, chunks
+ extra = text_chunks[0]["choices"][0]["delta"].get("extra_content")
+ assert extra == {"google": {"thought_signature": "SIG-TEXT"}}, text_chunks
+
+
+def test_openai_tools_translated_into_function_declarations(monkeypatch):
+ """Standard ChatCompletionRequest.tools must be forwarded into
+ Gemini's tools[].functionDeclarations envelope."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Look up the weather for a city.",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "city": {"type": "string"},
+ },
+ "required": ["city"],
+ },
+ },
+ }
+ ],
+ tool_choice = {"type": "function", "function": {"name": "get_weather"}},
+ )
+ tools_arr = captured["body"].get("tools") or []
+ fn_decls = [t for t in tools_arr if "functionDeclarations" in t]
+ assert fn_decls, captured["body"]
+ decls = fn_decls[0]["functionDeclarations"]
+ assert decls[0]["name"] == "get_weather"
+ assert decls[0]["parameters"]["properties"]["city"]["type"] == "string"
+ tool_config = captured["body"].get("toolConfig")
+ assert tool_config is not None, captured["body"]
+ fcc = tool_config["functionCallingConfig"]
+ assert fcc["mode"] == "ANY"
+ assert fcc["allowedFunctionNames"] == ["get_weather"]
+
+
+def test_tool_choice_auto_maps_to_function_calling_mode_auto(monkeypatch):
+ """tool_choice="auto" maps to toolConfig.functionCallingConfig.mode."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {"name": "noop", "parameters": {"type": "object"}},
+ }
+ ],
+ tool_choice = "auto",
+ )
+ fcc = captured["body"]["toolConfig"]["functionCallingConfig"]
+ assert fcc["mode"] == "AUTO"
+ assert "allowedFunctionNames" not in fcc
+
+
+def test_code_exec_inline_image_attaches_to_code_execution_card(monkeypatch):
+ """A codeExecution sandbox plot (matplotlib) ships as an inline
+ image part right after the codeExecutionResult. Instead of spawning
+ a separate empty image_generation card, attach to the same
+ code_execution tool_end via the `__IMAGES__:` marker the chat
+ adapter already understands."""
+ sse = [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "executableCode": {
+ "id": "code_plot",
+ "language": "PYTHON",
+ "code": "import matplotlib.pyplot as plt; plt.plot([1,2,3]); plt.savefig('out.png')",
+ },
+ },
+ {
+ "codeExecutionResult": {
+ "outcome": "OUTCOME_OK",
+ "output": "saved",
+ },
+ },
+ {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": base64.b64encode(b"PNGDATA").decode(),
+ },
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 5,
+ "candidatesTokenCount": 4,
+ },
+ }
+ ]
+ lines = _collect(
+ monkeypatch,
+ sse,
+ enabled_tools = ["code_execution"],
+ )
+ chunks = _parse_chunks(lines)
+ tool_events = [c["_toolEvent"] for c in chunks if "_toolEvent" in c]
+ # No standalone image_generation card should have been emitted.
+ image_starts = [
+ e
+ for e in tool_events
+ if e.get("type") == "tool_start" and e.get("tool_name") == "image_generation"
+ ]
+ assert not image_starts, tool_events
+ # The code_execution tool_end should now carry the inline image
+ # via the `__IMAGES__:` marker.
+ code_ends = [
+ e
+ for e in tool_events
+ if e.get("type") == "tool_end" and e.get("tool_call_id") == "code_plot"
+ ]
+ assert code_ends, tool_events
+ final_result = code_ends[-1]["result"]
+ assert "__IMAGES__:" in final_result, code_ends
+ assert "data:image/png;base64," in final_result, code_ends
+
+
+def test_code_execution_tool_call_replays_native_executable_code(monkeypatch):
+ """An assistant tool_call with toolName=code_execution and
+ extra_content.google.native_part containing the originally-emitted
+ `executableCode` + `codeExecutionResult` must round-trip as native
+ Gemini parts (not a generic functionCall) on the next turn."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "user", "content": "compute 2+2"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "code_a",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": "{}",
+ },
+ "extra_content": {
+ "google": {
+ "native_part": {
+ "executableCode": {
+ "id": "code_a",
+ "language": "PYTHON",
+ "code": "print(2+2)",
+ },
+ "codeExecutionResult": {
+ "outcome": "OUTCOME_OK",
+ "output": "4\n",
+ },
+ "thoughtSignature": "SIG-CODE",
+ },
+ },
+ },
+ },
+ ],
+ },
+ {"role": "user", "content": "what was that result"},
+ ],
+ )
+ assistant_turn = captured["body"]["contents"][1]
+ assert assistant_turn["role"] == "model"
+ parts = assistant_turn["parts"]
+ native_keys = [list(p.keys())[0] for p in parts if isinstance(p, dict)]
+ assert "executableCode" in native_keys, parts
+ assert "codeExecutionResult" in native_keys, parts
+ assert not any(
+ "functionCall" in p
+ and (p["functionCall"] or {}).get("name") == "code_execution"
+ for p in parts
+ ), parts
+ exec_part = next(p for p in parts if "executableCode" in p)
+ assert exec_part.get("thoughtSignature") == "SIG-CODE", exec_part
+
+
+def test_image_generation_tool_call_replays_native_inline_data(monkeypatch):
+ """An assistant tool_call with toolName=image_generation and
+ extra_content.google.native_part.inlineData must replay the prior
+ image as a native Gemini inlineData part (not a generic
+ functionCall) so multi-turn image editing keeps the image
+ context."""
+ pixel = base64.b64encode(b"PNG").decode()
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ messages = [
+ {"role": "user", "content": "make a circle"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "img_a",
+ "type": "function",
+ "function": {
+ "name": "image_generation",
+ "arguments": "{}",
+ },
+ "extra_content": {
+ "google": {
+ "native_part": {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": pixel,
+ },
+ "thoughtSignature": "SIG-IMG",
+ },
+ },
+ },
+ },
+ ],
+ },
+ {"role": "user", "content": "now make it blue"},
+ ],
+ )
+ assistant_turn = captured["body"]["contents"][1]
+ assert assistant_turn["role"] == "model"
+ parts = assistant_turn["parts"]
+ inline_parts = [p for p in parts if "inlineData" in p]
+ assert inline_parts, parts
+ assert inline_parts[0]["inlineData"]["mimeType"] == "image/png"
+ assert inline_parts[0]["inlineData"]["data"] == pixel
+ assert inline_parts[0].get("thoughtSignature") == "SIG-IMG", inline_parts
+ assert not any(
+ "functionCall" in p
+ and (p["functionCall"] or {}).get("name") == "image_generation"
+ for p in parts
+ ), parts
+
+
+def test_assistant_text_thought_signature_replays_on_outbound_text_part(monkeypatch):
+ """Assistant text with extra_content.google.thought_signature must
+ attach `thoughtSignature` to the LAST text part of the replayed
+ Gemini history. Gemini 3 strict function-calling rejects history
+ that drops returned signatures, so the frontend stows the latest
+ signed-text signature and the backend pins it on the next turn."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "text", "text": "hello"},
+ ],
+ "extra_content": {
+ "google": {"thought_signature": "SIG-TEXT"},
+ },
+ },
+ {"role": "user", "content": "again"},
+ ],
+ )
+ assistant_turn = captured["body"]["contents"][1]
+ assert assistant_turn["role"] == "model"
+ parts = assistant_turn["parts"]
+ text_parts = [p for p in parts if "text" in p]
+ assert text_parts, parts
+ assert text_parts[-1].get("thoughtSignature") == "SIG-TEXT", text_parts
+
+
+def test_function_declarations_strip_openai_only_schema_keys(monkeypatch):
+ """OpenAI strict tools commonly include `additionalProperties`,
+ `$schema`, `$defs`, `strict`, etc. Gemini's Schema rejects those
+ with INVALID_ARGUMENT, so the translator must strip them while
+ keeping properties..type intact."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "description": "Look up a value.",
+ "parameters": {
+ "type": "object",
+ "$schema": "http://json-schema.org/draft-07/schema#",
+ "additionalProperties": False,
+ "strict": True,
+ "properties": {
+ "key": {
+ "type": "string",
+ "additionalProperties": False,
+ },
+ },
+ "required": ["key"],
+ },
+ },
+ }
+ ],
+ )
+ tools_arr = captured["body"].get("tools") or []
+ decls = next(
+ (
+ t.get("functionDeclarations")
+ for t in tools_arr
+ if "functionDeclarations" in t
+ ),
+ None,
+ )
+ assert decls is not None, captured["body"]
+ params = decls[0]["parameters"]
+ assert "additionalProperties" not in params
+ assert "$schema" not in params
+ assert "strict" not in params
+ assert params["type"] == "object"
+ assert params["properties"]["key"]["type"] == "string"
+ assert "additionalProperties" not in params["properties"]["key"]
+ assert params["required"] == ["key"]
+
+
+def test_function_declarations_inline_local_refs_into_gemini_schema(monkeypatch):
+ """Round 25: Pydantic-generated tool schemas hoist nested object
+ shapes into `$defs` and reference them with `{"$ref": "#/$defs/..."}`.
+ Gemini's OpenAPI subset has no $ref, so a naive allowlist sanitizer
+ drops the reference and reduces the nested property to `{}`, losing
+ its type, fields, and required keys. The sanitizer must resolve
+ local `#/...` pointers and inline the referenced schema."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "set_user",
+ "description": "Persist a user.",
+ "parameters": {
+ "type": "object",
+ "$defs": {
+ "Address": {
+ "type": "object",
+ "properties": {
+ "street": {"type": "string"},
+ "zip": {"type": "string"},
+ },
+ "required": ["street", "zip"],
+ },
+ },
+ "properties": {
+ "name": {"type": "string"},
+ "address": {"$ref": "#/$defs/Address"},
+ },
+ "required": ["name", "address"],
+ },
+ },
+ }
+ ],
+ )
+ tools_arr = captured["body"].get("tools") or []
+ decls = next(
+ (
+ t.get("functionDeclarations")
+ for t in tools_arr
+ if "functionDeclarations" in t
+ ),
+ None,
+ )
+ assert decls is not None, captured["body"]
+ params = decls[0]["parameters"]
+ assert "$defs" not in params
+ address = params["properties"]["address"]
+ assert address.get("type") == "object", address
+ assert address.get("properties", {}).get("street", {}).get("type") == "string"
+ assert address.get("properties", {}).get("zip", {}).get("type") == "string"
+ assert address.get("required") == ["street", "zip"]
+
+
+def test_function_declarations_inline_local_refs_in_anyof_and_items(monkeypatch):
+ """The recursive inliner must reach through `anyOf` branches and
+ `items` (array element schemas) as well, not just top-level
+ property refs."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "bulk_set",
+ "parameters": {
+ "type": "object",
+ "$defs": {
+ "Address": {
+ "type": "object",
+ "properties": {"zip": {"type": "string"}},
+ "required": ["zip"],
+ },
+ },
+ "properties": {
+ "primary": {
+ "anyOf": [
+ {"$ref": "#/$defs/Address"},
+ {"type": "null"},
+ ],
+ },
+ "extras": {
+ "type": "array",
+ "items": {"$ref": "#/$defs/Address"},
+ },
+ },
+ },
+ },
+ }
+ ],
+ )
+ tools_arr = captured["body"].get("tools") or []
+ decls = next(
+ (
+ t.get("functionDeclarations")
+ for t in tools_arr
+ if "functionDeclarations" in t
+ ),
+ None,
+ )
+ assert decls is not None
+ params = decls[0]["parameters"]
+ primary = params["properties"]["primary"]
+ # anyOf with single non-null branch + null collapses to inline +
+ # nullable: true, and the inlined branch must contain the resolved
+ # Address shape.
+ assert primary.get("nullable") is True
+ assert primary.get("type") == "object"
+ assert primary.get("properties", {}).get("zip", {}).get("type") == "string"
+ extras = params["properties"]["extras"]
+ assert extras.get("type") == "array"
+ assert extras.get("items", {}).get("type") == "object"
+ assert (
+ extras.get("items", {}).get("properties", {}).get("zip", {}).get("type")
+ == "string"
+ )
+
+
+def test_function_declarations_self_referential_schema_terminates(monkeypatch):
+ """Self-referential / cyclic JSON Schemas (a `Node` that contains
+ `children: [Node]`) must not infinite-loop. The inliner tracks the
+ set of refs in flight and short-circuits to `{}` on a cycle."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "set_tree",
+ "parameters": {
+ "type": "object",
+ "$defs": {
+ "Node": {
+ "type": "object",
+ "properties": {
+ "value": {"type": "string"},
+ "children": {
+ "type": "array",
+ "items": {"$ref": "#/$defs/Node"},
+ },
+ },
+ },
+ },
+ "properties": {
+ "root": {"$ref": "#/$defs/Node"},
+ },
+ },
+ },
+ }
+ ],
+ )
+ tools_arr = captured["body"].get("tools") or []
+ decls = next(
+ (
+ t.get("functionDeclarations")
+ for t in tools_arr
+ if "functionDeclarations" in t
+ ),
+ None,
+ )
+ assert decls is not None
+ root = decls[0]["parameters"]["properties"]["root"]
+ assert root.get("type") == "object"
+ assert root.get("properties", {}).get("value", {}).get("type") == "string"
+
+
+def test_gemini_native_skips_orphan_function_response_for_dropped_builtin(
+ monkeypatch,
+):
+ """Round 26: when the assistant-side synthetic web_search/web_fetch
+ tool_call is dropped from native Gemini history, the matching
+ role="tool" follow-up must also be dropped. Otherwise the outbound
+ body carries an orphan functionResponse with no preceding
+ functionCall, which 400s the Gemini turn."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ req = ChatCompletionRequest.model_validate(
+ {
+ "model": "gemini-2.5-flash",
+ "messages": [
+ {"role": "user", "content": "search please"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_s",
+ "type": "function",
+ "function": {
+ "name": "web_search",
+ "arguments": ('{"_server_tool": true, "query": "x"}'),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_s",
+ "content": "[search result]",
+ },
+ {"role": "user", "content": "again"},
+ ],
+ "max_tokens": 64,
+ "stream": True,
+ }
+ )
+ built = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta",
+ )
+ captured = _capture_body(monkeypatch, messages = built)
+ contents = captured["body"].get("contents") or []
+ for entry in contents:
+ for part in entry.get("parts", []):
+ fr = part.get("functionResponse")
+ if isinstance(fr, dict):
+ assert fr.get("name") != "web_search", contents
+
+
+def test_gemini_native_skips_orphan_function_response_for_native_part_replay(
+ monkeypatch,
+):
+ """Round 26: code_execution / image_generation tool_calls are
+ replayed as Gemini-native executableCode / codeExecutionResult /
+ inlineData parts. The matching role="tool" follow-up must NOT then
+ be emitted as a functionResponse named code_execution -- there is
+ no declared user function with that name, and Gemini's history
+ rules already attribute the result to the native parts above."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "user", "content": "plot something"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_a",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": "{}",
+ },
+ "extra_content": {
+ "google": {
+ "native_part": {
+ "parts": [
+ {
+ "executableCode": {
+ "language": "PYTHON",
+ "code": "print(2)",
+ }
+ },
+ {
+ "codeExecutionResult": {
+ "outcome": "OUTCOME_OK",
+ "output": "2\n",
+ }
+ },
+ ]
+ }
+ }
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_a",
+ "name": "code_execution",
+ "content": "2",
+ },
+ {"role": "user", "content": "next"},
+ ],
+ )
+ contents = captured["body"].get("contents") or []
+ saw_native = False
+ for entry in contents:
+ for part in entry.get("parts", []):
+ if "executableCode" in part or "codeExecutionResult" in part:
+ saw_native = True
+ fr = part.get("functionResponse")
+ if isinstance(fr, dict):
+ assert fr.get("name") != "code_execution", contents
+ assert saw_native, contents
+
+
+def test_gemini_native_part_falls_back_to_args_google(monkeypatch):
+ """Round 27: a direct OpenAI-compat API caller (or imported third-
+ party thread) cannot use Studio's non-standard
+ `tool_calls[].extra_content` field, so the native_part payload
+ round-trips through `function.arguments` as
+ `{"google": {"native_part": {...}}}`. The synthetic-builtin
+ detector recognizes that location, but the replay branch was only
+ reading from `tc.extra_content.google.native_part`. Result: the
+ round-25 guard saw a synthetic builtin with no _native_part and
+ dropped the entire assistant turn, losing the prior code/image
+ context. The translator must fall back to args.google.native_part
+ and still emit the native executableCode / inlineData parts."""
+ import json as _json
+
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "user", "content": "draw a cat"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_img",
+ "type": "function",
+ "function": {
+ "name": "image_generation",
+ "arguments": _json.dumps(
+ {
+ "google": {
+ "native_part": {
+ "parts": [
+ {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": "AAAA",
+ }
+ }
+ ]
+ }
+ }
+ }
+ ),
+ },
+ }
+ ],
+ },
+ {"role": "user", "content": "now make it a dog"},
+ ],
+ )
+ contents = captured["body"].get("contents") or []
+ saw_inline = False
+ for entry in contents:
+ for part in entry.get("parts", []):
+ if "inlineData" in part:
+ saw_inline = True
+ assert saw_inline, contents
+
+
+def test_gemini_native_skips_synthetic_server_builtin_replay(monkeypatch):
+ """Round 25: Marked server-side builtin tool_calls (web_search /
+ web_fetch with `_server_tool` or `args.google.native_part`) must
+ not fall through to the generic Gemini `functionCall` replay path
+ when no replayable native part exists. Without this guard the
+ outbound body contains a fake `functionCall` whose name is not a
+ declared user function, and the Gemini turn 400s."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ req = ChatCompletionRequest.model_validate(
+ {
+ "model": "gemini-2.5-flash",
+ "messages": [
+ {"role": "user", "content": "search please"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_s",
+ "type": "function",
+ "function": {
+ "name": "web_search",
+ "arguments": ('{"_server_tool": true, "query": "x"}'),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_s",
+ "content": "[search result]",
+ },
+ {"role": "user", "content": "again"},
+ ],
+ "max_tokens": 64,
+ "stream": True,
+ }
+ )
+ built = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta",
+ )
+ captured = _capture_body(monkeypatch, messages = built)
+ contents = captured["body"].get("contents") or []
+ for entry in contents:
+ for part in entry.get("parts", []):
+ fc = part.get("functionCall")
+ if isinstance(fc, dict):
+ assert fc.get("name") != "web_search", contents
+
+
+def test_chat_message_extra_content_round_trips_through_validation():
+ """Round 9: ChatMessage was missing `extra_content`, so Pydantic
+ discarded the field during request validation and the text-part
+ signature replay path read nothing. The field must survive
+ model_validate and pass through _build_external_messages."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ req = ChatCompletionRequest.model_validate(
+ {
+ "model": "gemini-2.5-flash",
+ "messages": [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": [
+ {"type": "text", "text": "hello"},
+ ],
+ "extra_content": {
+ "google": {"thought_signature": "SIG-TEXT"},
+ },
+ },
+ {"role": "user", "content": "again"},
+ ],
+ "max_tokens": 64,
+ "stream": True,
+ }
+ )
+ assistant_msg = req.messages[1]
+ assert assistant_msg.extra_content == {
+ "google": {"thought_signature": "SIG-TEXT"},
+ }
+ built = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta",
+ )
+ assistant_out = built[1]
+ assert assistant_out["extra_content"] == {
+ "google": {"thought_signature": "SIG-TEXT"},
+ }
+ # Non-Gemini providers must NOT receive extra_content; Google's
+ # thought_signature field is unknown to OpenAI / Mistral / etc.
+ built_openai = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "openai",
+ )
+ assert "extra_content" not in built_openai[1], built_openai[1]
+ # Custom non-Google Gemini bases (LiteLLM / OAI-compat gateways)
+ # also must not receive Gemini-only extra_content because the
+ # backend dispatches them through /chat/completions.
+ built_custom = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "gemini",
+ base_url = "https://litellm.example/v1",
+ )
+ assert "extra_content" not in built_custom[1], built_custom[1]
+
+
+def test_parallel_tool_results_group_into_one_user_block(monkeypatch):
+ """Round 14: Gemini docs show parallel functionResponses grouped
+ in a single subsequent user content with multiple
+ functionResponse parts. Consecutive OpenAI role="tool" messages
+ must merge into one Gemini user block, not split into separate
+ user turns."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {"role": "user", "content": "compute"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "call_a",
+ "type": "function",
+ "function": {"name": "add", "arguments": '{"x":1}'},
+ },
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {"name": "mul", "arguments": '{"x":2}'},
+ },
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_a",
+ "name": "add",
+ "content": "2",
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_b",
+ "name": "mul",
+ "content": "4",
+ },
+ ],
+ )
+ contents = captured["body"]["contents"]
+ # Initial user, model with two functionCalls, ONE user with two
+ # functionResponses.
+ tool_result_users = [
+ c
+ for c in contents
+ if c.get("role") == "user"
+ and all(
+ isinstance(p, dict) and "functionResponse" in p
+ for p in (c.get("parts") or [])
+ )
+ ]
+ assert len(tool_result_users) == 1, contents
+ fr_parts = tool_result_users[0]["parts"]
+ assert len(fr_parts) == 2, fr_parts
+ names = [p["functionResponse"]["name"] for p in fr_parts]
+ assert names == ["add", "mul"], names
+
+
+def test_function_schema_nullable_type_array_flattens(monkeypatch):
+ """Round 14: OpenAI strict tools commonly use
+ `"type": ["string", "null"]` for optional fields. Gemini's
+ OpenAPI-style Schema rejects union types and expects
+ `"type": "string"` with `"nullable": true`. The sanitizer must
+ translate the union form."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "city": {"type": ["string", "null"]},
+ "score": {"type": ["number", "null"]},
+ },
+ },
+ },
+ }
+ ],
+ )
+ decls = next(
+ t["functionDeclarations"]
+ for t in captured["body"].get("tools") or []
+ if "functionDeclarations" in t
+ )
+ params = decls[0]["parameters"]["properties"]
+ assert params["city"]["type"] == "string"
+ assert params["city"]["nullable"] is True
+ assert params["score"]["type"] == "number"
+ assert params["score"]["nullable"] is True
+
+
+def test_image_picker_model_with_search_off_pill_strips_text_tools(monkeypatch):
+ """Round 11: image-tier model id rejects text-only tools and
+ thinkingConfig at the model level regardless of whether the Images
+ pill is on. Selecting gemini-2.5-flash-image + enabled_tools=
+ ["web_search"] with no image_generation must NOT forward
+ googleSearch or thinkingConfig (Gemini 400s on text tools for
+ legacy image ids)."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["web_search"],
+ reasoning_effort = "high",
+ )
+ body = captured["body"]
+ assert "tools" not in body, body.get("tools")
+ assert "thinkingConfig" not in body.get("generationConfig", {}), body[
+ "generationConfig"
+ ]
+
+
+def test_image_models_drop_function_declarations(monkeypatch):
+ """Image-mode requests cannot mix tools with responseModalities so
+ user-supplied function declarations must be dropped."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation"],
+ tools = [
+ {
+ "type": "function",
+ "function": {"name": "noop", "parameters": {"type": "object"}},
+ }
+ ],
+ )
+ assert captured["body"].get("tools") is None
+ assert captured["body"]["generationConfig"]["responseModalities"] == [
+ "TEXT",
+ "IMAGE",
+ ]
+
+
+def test_safe_fetch_image_rejects_malformed_bracketed_url():
+ """Round 17: bracketed IPv6 garbage like `https://[bad/x.png` makes
+ urlparse raise ValueError. The fetch helper must catch it and drop
+ the image rather than crashing the request mid-build."""
+ res = _drive(ep_mod._safe_fetch_image_for_gemini("https://[bad/x.png", "image/png"))
+ assert res is None
+
+
+def test_safe_fetch_image_pins_validated_ip_no_hostname_in_request(
+ monkeypatch,
+):
+ """Round 17: the fetch helper must pin the validated IP into the
+ outgoing request URL (with a Host header carrying the original
+ hostname). A second hostname-style getaddrinfo after the validate
+ step would be a DNS-rebinding gap, so we assert the urllib opener
+ is called with an IP-rewritten URL."""
+ import socket
+
+ captured: dict = {"requests": []}
+
+ # Public IP during validate; record every getaddrinfo call.
+ original_getaddrinfo = socket.getaddrinfo
+
+ def fake_getaddrinfo(host, *args, **kwargs):
+ captured.setdefault("dns", []).append(host)
+ if host == "cdn.example.com":
+ return [
+ (
+ socket.AF_INET,
+ socket.SOCK_STREAM,
+ 0,
+ "",
+ ("8.8.8.8", 0),
+ )
+ ]
+ return original_getaddrinfo(host, *args, **kwargs)
+
+ monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
+
+ class _StubResp:
+ status = 200
+ headers = {"content-type": "image/png", "content-length": "3"}
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *a):
+ return False
+
+ def read(self, _n = None):
+ return b"PNG"
+
+ class _StubOpener:
+ def open(self, req, timeout = None):
+ captured["requests"].append(
+ {
+ "url": req.full_url,
+ "host_header": req.get_header("Host"),
+ }
+ )
+ return _StubResp()
+
+ monkeypatch.setattr(
+ "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()
+ )
+
+ res = _drive(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://cdn.example.com/x.png", "image/png"
+ )
+ )
+ assert res is not None
+ assert res[0] == "image/png"
+ # The outgoing URL must use the pinned IP literal, not the hostname.
+ assert any("8.8.8.8" in r["url"] for r in captured["requests"]), captured
+ assert all(
+ "cdn.example.com" not in r["url"] for r in captured["requests"]
+ ), captured
+ # Host header still carries the original hostname for vhost/SNI.
+ assert captured["requests"][0]["host_header"] == "cdn.example.com"
+
+
+def test_safe_fetch_image_redirect_to_private_host_rejected(monkeypatch):
+ """Round 17: each redirect hop must re-validate the new host. A
+ public hop that redirects to an internal address must be dropped."""
+ import socket
+ import urllib.error
+
+ original_getaddrinfo = socket.getaddrinfo
+
+ def fake_getaddrinfo(host, *args, **kwargs):
+ if host == "cdn.example.com":
+ return [
+ (
+ socket.AF_INET,
+ socket.SOCK_STREAM,
+ 0,
+ "",
+ ("1.1.1.1", 0),
+ )
+ ]
+ if host == "internal.bad":
+ return [
+ (
+ socket.AF_INET,
+ socket.SOCK_STREAM,
+ 0,
+ "",
+ ("10.0.0.5", 0),
+ )
+ ]
+ return original_getaddrinfo(host, *args, **kwargs)
+
+ monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
+
+ class _StubOpener:
+ def open(self, req, timeout = None):
+ # Simulate a 302 to a private host.
+ raise urllib.error.HTTPError(
+ req.full_url,
+ 302,
+ "Found",
+ {"Location": "https://internal.bad/secret.png"},
+ None,
+ )
+
+ monkeypatch.setattr(
+ "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()
+ )
+
+ res = _drive(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://cdn.example.com/x.png", "image/png"
+ )
+ )
+ assert res is None
+
+
+def test_files_api_substring_url_not_misclassified_as_filedata(monkeypatch):
+ """Round 17: a CDN URL whose path/query merely contains the Files
+ API substring must NOT be sent as `fileData.fileUri`; it must be
+ routed through the safe-fetch path. Previously the substring check
+ `"generativelanguage.googleapis.com/" in url.lower()` matched any
+ URL carrying that text anywhere."""
+ captured_outbound: dict = {}
+ fetch_calls: list[str] = []
+
+ async def fake_fetch(url, fallback_mime, max_bytes = None):
+ fetch_calls.append(url)
+ return "image/png", base64.b64encode(b"DATA").decode("ascii")
+
+ monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch)
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured_outbound["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = _gemini_sse(
+ [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "ok"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 1,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ ),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = _make_gemini_client()
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "describe"},
+ {
+ "type": "image_url",
+ "image_url": {
+ # Looks like a Files API URL in the path
+ # but the host is an attacker CDN.
+ "url": "https://evil.example/path/generativelanguage.googleapis.com/v1beta/files/abc.png",
+ },
+ },
+ {
+ "type": "image_url",
+ "image_url": {
+ # Looks YouTube-ish in the path.
+ "url": "https://cdn.example.com/youtube.com/cat.png",
+ },
+ },
+ ],
+ }
+ ],
+ model = "gemini-2.5-flash",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 64,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ parts = captured_outbound["body"]["contents"][-1]["parts"]
+ assert not any("fileData" in p for p in parts), parts
+ inline_count = sum(1 for p in parts if "inlineData" in p)
+ assert inline_count == 2, parts
+ assert len(fetch_calls) == 2, fetch_calls
+
+
+def test_function_schema_anyof_null_variant_flattens_to_nullable(monkeypatch):
+ """Round 17: OpenAI/Pydantic emit `anyOf: [{X}, {"type":"null"}]`
+ for Optional[X]. Gemini's OpenAPI subset rejects `"type":"null"`
+ inside anyOf. The sanitizer must collapse a singleton-plus-null
+ union back to the non-null branch with `nullable: true`."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "label": {
+ "anyOf": [
+ {"type": "string"},
+ {"type": "null"},
+ ]
+ },
+ "count": {
+ "anyOf": [
+ {"type": "integer"},
+ {"type": "null"},
+ ]
+ },
+ },
+ },
+ },
+ }
+ ],
+ )
+ decls = next(
+ t["functionDeclarations"]
+ for t in captured["body"].get("tools") or []
+ if "functionDeclarations" in t
+ )
+ params = decls[0]["parameters"]["properties"]
+ assert params["label"]["type"] == "string"
+ assert params["label"]["nullable"] is True
+ assert "anyOf" not in params["label"]
+ assert params["count"]["type"] == "integer"
+ assert params["count"]["nullable"] is True
+
+
+def test_legacy_gemini3_pro_medium_coerced_to_high(monkeypatch):
+ """Round 17: legacy `gemini-3-pro*` (including `-preview`, shut down
+ 2026-03-09) only accepted low/high. 3.1+ Pro added medium. The
+ backend must coerce medium → high for the legacy model so stale UI
+ state does not 400 the request."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3-pro-preview",
+ reasoning_effort = "medium",
+ )
+ assert captured["body"]["generationConfig"]["thinkingConfig"] == {
+ "thinkingLevel": "high",
+ }
+
+
+def test_gemini_3_1_pro_medium_passes_through(monkeypatch):
+ """Round 17 regression: 3.1+ Pro accepts medium; coercion must NOT
+ apply when the model id is gemini-3.1-pro*."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-3.1-pro-preview",
+ reasoning_effort = "medium",
+ )
+ assert captured["body"]["generationConfig"]["thinkingConfig"] == {
+ "thinkingLevel": "medium",
+ }
+
+
+def test_tool_calls_extra_content_stripped_for_non_native_gemini():
+ """Round 17: per-tool-call `extra_content` (Gemini thoughtSignature
+ carrier) must not leak through `_build_external_messages` to
+ non-native-Gemini providers; OpenAI / Anthropic / custom Gemini
+ OAI-compat gateways would 400 on the unknown key."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_1",
+ "type": "function",
+ "function": {"name": "lookup", "arguments": "{}"},
+ "extra_content": {
+ "google": {"thought_signature": "SIG"},
+ },
+ }
+ ],
+ }
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+
+ # Non-native providers (openai, custom Gemini OAI-compat proxy)
+ # must have extra_content stripped from the tool_call entry.
+ for provider_type, base_url in [
+ ("openai", None),
+ ("gemini", "https://litellm.example/v1"),
+ ]:
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = provider_type,
+ base_url = base_url,
+ )
+ assert len(result) == 1
+ tc = result[0]["tool_calls"][0]
+ assert "extra_content" not in tc, (provider_type, tc)
+
+ # Native Gemini still receives extra_content for the round-trip.
+ result_native = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta",
+ )
+ tc_native = result_native[0]["tool_calls"][0]
+ assert tc_native["extra_content"]["google"]["thought_signature"] == "SIG"
+
+
+def test_user_function_named_with_server_tool_arg_not_dropped(monkeypatch):
+ """Round 17: the OpenAI Responses translator must NOT drop a user
+ function whose JSON arguments happen to contain `_server_tool:
+ true` UNLESS the function name is also one of the canonical
+ builtin names. Otherwise a user schema with an `_server_tool` field
+ becomes invisible to the model."""
+ captured: dict = {"input_items": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ body = json.loads(request.content.decode("utf-8"))
+ captured["input_items"] = body.get("input")
+ return httpx.Response(
+ 200,
+ content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openai",
+ base_url = "https://api.openai.com/v1",
+ api_key = "sk-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_user",
+ "type": "function",
+ "function": {
+ "name": "user_function",
+ "arguments": json.dumps(
+ {"_server_tool": True, "q": "x"}
+ ),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "content": "result",
+ "tool_call_id": "call_user",
+ "name": "user_function",
+ },
+ {"role": "user", "content": "continue"},
+ ],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 1.0,
+ max_tokens = 16,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ items = captured["input_items"] or []
+ fn_calls = [i for i in items if i.get("type") == "function_call"]
+ fn_outs = [i for i in items if i.get("type") == "function_call_output"]
+ # User function call must survive (matching call + output).
+ assert any(c.get("name") == "user_function" for c in fn_calls), items
+ assert len(fn_outs) == 1, items
+
+
+def test_builtin_named_with_server_tool_marker_dropped(monkeypatch):
+ """Round 17 control: a builtin (web_search) tagged with
+ `_server_tool: true` continues to be filtered from outbound
+ history."""
+ captured: dict = {"input_items": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ body = json.loads(request.content.decode("utf-8"))
+ captured["input_items"] = body.get("input")
+ return httpx.Response(
+ 200,
+ content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openai",
+ base_url = "https://api.openai.com/v1",
+ api_key = "sk-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {"role": "user", "content": "search please"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {
+ "name": "web_search",
+ "arguments": json.dumps(
+ {"_server_tool": True, "query": "x"}
+ ),
+ },
+ }
+ ],
+ },
+ {"role": "user", "content": "continue"},
+ ],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 1.0,
+ max_tokens = 16,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ items = captured["input_items"] or []
+ fn_calls = [i for i in items if i.get("type") == "function_call"]
+ # Builtin server-side tool call must be filtered out.
+ assert all(c.get("name") != "web_search" for c in fn_calls), items
+
+
+def test_gemini_tool_choice_none_disables_hosted_builtins(monkeypatch):
+ """Round 18: `tool_choice="none"` must drop hosted Google Search /
+ code execution from the outbound Gemini body, not just user
+ function declarations. Otherwise an API client that opted out of
+ tool use still triggers grounded search (privacy + billing)."""
+ captured = _capture_body(
+ monkeypatch,
+ enabled_tools = ["web_search", "code_execution"],
+ tool_choice = "none",
+ )
+ assert captured["body"].get("tools") is None, captured["body"]
+
+
+def test_gemini_tool_choice_none_disables_function_declarations(monkeypatch):
+ """Round 18: `tool_choice="none"` must drop user function
+ declarations as well as hosted builtins from the Gemini body."""
+ captured = _capture_body(
+ monkeypatch,
+ tool_choice = "none",
+ tools = [
+ {
+ "type": "function",
+ "function": {"name": "lookup", "parameters": {"type": "object"}},
+ }
+ ],
+ )
+ assert captured["body"].get("tools") is None, captured["body"]
+
+
+def test_schema_anyof_multitype_with_null_keeps_anyof_and_nullable(
+ monkeypatch,
+):
+ """Round 18: multi-branch unions with null (e.g.
+ `Union[str, int, None]`) must keep the slim anyOf without the null
+ branch and add `nullable: true`; Gemini rejects
+ `{"type":"null"}` inside anyOf."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "either": {
+ "anyOf": [
+ {"type": "string"},
+ {"type": "integer"},
+ {"type": "null"},
+ ]
+ },
+ },
+ },
+ },
+ }
+ ],
+ )
+ decls = next(
+ t["functionDeclarations"]
+ for t in captured["body"].get("tools") or []
+ if "functionDeclarations" in t
+ )
+ either = decls[0]["parameters"]["properties"]["either"]
+ assert either.get("nullable") is True
+ inner = either.get("anyOf")
+ assert isinstance(inner, list) and len(inner) == 2, either
+ assert all(
+ not (isinstance(b, dict) and b.get("type") == "null") for b in inner
+ ), inner
+
+
+def test_safe_fetch_image_redirect_malformed_url_no_crash(monkeypatch):
+ """Round 18: when the upstream 302 Location is a malformed
+ bracketed-IPv6 URL, the helper must return None instead of letting
+ a urlparse ValueError abort the chat stream."""
+ import socket
+ import urllib.error
+
+ original_getaddrinfo = socket.getaddrinfo
+
+ def fake_getaddrinfo(host, *args, **kwargs):
+ if host == "cdn.example.com":
+ return [
+ (
+ socket.AF_INET,
+ socket.SOCK_STREAM,
+ 0,
+ "",
+ ("1.1.1.1", 0),
+ )
+ ]
+ return original_getaddrinfo(host, *args, **kwargs)
+
+ monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
+
+ class _StubOpener:
+ def open(self, req, timeout = None):
+ raise urllib.error.HTTPError(
+ req.full_url,
+ 302,
+ "Found",
+ {"Location": "https://[bad/x.png"},
+ None,
+ )
+
+ monkeypatch.setattr(
+ "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()
+ )
+
+ res = _drive(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://cdn.example.com/x.png", "image/png"
+ )
+ )
+ assert res is None
+
+
+def test_safe_fetch_image_malformed_port_no_crash():
+ """Round 18: a URL with a non-numeric port (`https://h:bad/x.png`)
+ must not raise; urlparse's port property lazily ValueErrors."""
+ res = _drive(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://example.com:bad/x.png", "image/png"
+ )
+ )
+ assert res is None
+
+
+def test_safe_fetch_image_missing_content_type_uses_fallback(monkeypatch):
+ """Round 18: when the server returns image bytes but no
+ Content-Type header, the helper must use the caller-provided
+ fallback MIME (guessed from URL extension) instead of dropping the
+ image as `non-image content-type=`."""
+ import socket
+
+ original_getaddrinfo = socket.getaddrinfo
+
+ def fake_getaddrinfo(host, *args, **kwargs):
+ if host == "cdn.example.com":
+ return [
+ (
+ socket.AF_INET,
+ socket.SOCK_STREAM,
+ 0,
+ "",
+ ("1.1.1.1", 0),
+ )
+ ]
+ return original_getaddrinfo(host, *args, **kwargs)
+
+ monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
+
+ class _StubResp:
+ status = 200
+ headers = {"content-length": "3"}
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *a):
+ return False
+
+ def read(self, _n = None):
+ return b"PNG"
+
+ class _StubOpener:
+ def open(self, req, timeout = None):
+ return _StubResp()
+
+ monkeypatch.setattr(
+ "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()
+ )
+
+ res = _drive(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://cdn.example.com/cat.png", "image/png"
+ )
+ )
+ assert res is not None
+ assert res[0] == "image/png"
+
+
+def test_anthropic_translates_openai_tool_calls_into_tool_use_blocks(monkeypatch):
+ """Round 18: an assistant turn with OpenAI-style top-level
+ `tool_calls` must be translated into Anthropic native
+ `{type:"tool_use", id, name, input}` content blocks before being
+ forwarded. The OpenAI `role="tool"` follow-up must become a
+ `role:"user"` message with a `tool_result` content block."""
+ captured: dict = {"messages": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ body = json.loads(request.content.decode("utf-8"))
+ captured["messages"] = body.get("messages")
+ return httpx.Response(
+ 200,
+ content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "anthropic",
+ base_url = "https://api.anthropic.com",
+ api_key = "sk-ant-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {"role": "user", "content": "look up X"},
+ {
+ "role": "assistant",
+ "content": "let me check",
+ "tool_calls": [
+ {
+ "id": "call_a",
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "arguments": '{"q":"x"}',
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "content": "result_text",
+ "tool_call_id": "call_a",
+ "name": "lookup",
+ },
+ {"role": "user", "content": "summarise"},
+ ],
+ model = "claude-sonnet-4-5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 64,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ msgs = captured["messages"] or []
+ # No top-level tool_calls should remain.
+ assert all("tool_calls" not in m for m in msgs), msgs
+ # The assistant turn must now have content blocks including a
+ # tool_use block.
+ asst = [m for m in msgs if m.get("role") == "assistant"]
+ assert asst and isinstance(asst[0]["content"], list), asst
+ tool_uses = [b for b in asst[0]["content"] if b.get("type") == "tool_use"]
+ assert len(tool_uses) == 1, asst[0]
+ assert tool_uses[0]["name"] == "lookup"
+ assert tool_uses[0]["input"] == {"q": "x"}
+ # The role="tool" message must become a user/tool_result message.
+ tool_results: list[dict] = []
+ for m in msgs:
+ if m.get("role") == "user" and isinstance(m.get("content"), list):
+ tool_results.extend(
+ b for b in m["content"] if b.get("type") == "tool_result"
+ )
+ assert any(
+ tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text"
+ for tr in tool_results
+ ), msgs
+
+
+def test_unmarked_user_web_search_function_survives_serialization():
+ """Round 18: a user-defined function literally named `web_search`
+ with NO `_server_tool` marker must survive `_build_external_messages`
+ when forwarded to a non-native provider; only marked synthetic
+ builtin cards may be dropped."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_user",
+ "type": "function",
+ "function": {
+ "name": "web_search",
+ "arguments": '{"query": "x"}',
+ },
+ }
+ ],
+ }
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "openai",
+ base_url = None,
+ )
+ assert len(result) == 1, result
+ tcs = result[0].get("tool_calls") or []
+ assert len(tcs) == 1, result
+ assert tcs[0]["function"]["name"] == "web_search"
+
+
+def test_marked_server_builtin_dropped_from_build_external_messages():
+ """Round 18: when a Gemini-native turn carrying a marked
+ `image_generation` server-tool card is forwarded to OpenAI / a
+ custom Gemini OAI-compat proxy, the tool_call must be dropped, not
+ just have its extra_content stripped. Forwarding an orphan
+ `image_generation` tool_call would 400 the receiving API."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ marked_args = json.dumps({"_server_tool": True, "kind": "image"})
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {
+ "name": "image_generation",
+ "arguments": marked_args,
+ },
+ }
+ ],
+ }
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+ # Non-native providers: marked builtin tool_call must be dropped
+ # AND if it was the only payload, the whole message disappears.
+ for provider_type, base_url in [
+ ("openai", None),
+ ("gemini", "https://litellm.example/v1"),
+ ]:
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = provider_type,
+ base_url = base_url,
+ )
+ # Empty assistant turn with only synthetic tool_call dropped.
+ assert result == [] or all(not (m.get("tool_calls") or []) for m in result), (
+ provider_type,
+ result,
+ )
+
+ # Native Gemini preserves it (round-trips via extra_content).
+ result_native = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "gemini",
+ base_url = "https://generativelanguage.googleapis.com/v1beta",
+ )
+ assert len(result_native) == 1
+ assert result_native[0]["tool_calls"][0]["function"]["name"] == "image_generation"
+
+
+def test_openai_responses_tool_choice_none_drops_hosted_tools(monkeypatch):
+ """Round 18: `tool_choice="none"` must also drop hosted OpenAI
+ Responses builtins (web_search, code execution shell, image
+ generation), not just user function tools."""
+ captured: dict = {"body": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openai",
+ base_url = "https://api.openai.com/v1",
+ api_key = "sk-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 1.0,
+ max_tokens = 16,
+ enabled_tools = ["web_search", "code_execution", "image_generation"],
+ tool_choice = "none",
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ body = captured["body"] or {}
+ assert body.get("tools") in (None, []), body
+
+
+def test_anthropic_tool_choice_none_drops_hosted_tools(monkeypatch):
+ """Round 19: tool_choice="none" must opt out of Anthropic hosted
+ builtins (web_search, web_fetch, code_execution) just like it does
+ for Gemini and OpenAI Responses."""
+ captured: dict = {"body": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "anthropic",
+ base_url = "https://api.anthropic.com",
+ api_key = "sk-ant-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "claude-sonnet-4-5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search", "web_fetch", "code_execution"],
+ tool_choice = "none",
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ body = captured["body"] or {}
+ assert body.get("tools") in (None, []), body
+
+
+def test_openrouter_tool_choice_none_drops_web_plugin(monkeypatch):
+ """Round 19: tool_choice="none" must drop the OpenRouter web
+ plugin so a request that opted out of tool use does not still
+ trigger hosted web search."""
+ captured: dict = {"body": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = b"data: [DONE]\n\n",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openrouter",
+ base_url = "https://openrouter.ai/api/v1",
+ api_key = "sk-or-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "openai/gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search"],
+ tool_choice = "none",
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ body = captured["body"] or {}
+ assert body.get("plugins") in (None, []), body
+
+
+def test_kimi_tool_choice_none_skips_web_search_helper(monkeypatch):
+ """Round 19: when tool_choice="none" plus enabled_tools=
+ ["web_search"] on Kimi, the dispatcher must NOT route into
+ `_stream_kimi_web_search`. Falling through to the generic OAI-
+ compat path is the expected behavior."""
+ routed_to_helper = {"called": False}
+
+ real_helper = ExternalProviderClient._stream_kimi_web_search
+
+ async def fake_helper(self, *args, **kwargs): # noqa: ARG001
+ routed_to_helper["called"] = True
+ if False:
+ yield "" # pragma: no cover
+
+ monkeypatch.setattr(
+ ExternalProviderClient,
+ "_stream_kimi_web_search",
+ fake_helper,
+ )
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = b"data: [DONE]\n\n",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "kimi",
+ base_url = "https://api.moonshot.ai/v1",
+ api_key = "sk-kimi-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "kimi-k2.6",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search"],
+ tool_choice = "none",
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ assert routed_to_helper["called"] is False
+
+ monkeypatch.setattr(
+ ExternalProviderClient,
+ "_stream_kimi_web_search",
+ real_helper,
+ )
+
+
+def test_user_code_execution_function_not_dropped():
+ """Round 19: a user-declared function literally named
+ `code_execution` with normal `code` arguments must survive
+ `_build_external_messages` -- round 17's shape heuristic dropped
+ it, which broke function-calling round-trips."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_user",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": '{"code": "print(1)"}',
+ },
+ }
+ ],
+ }
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "openai",
+ base_url = None,
+ )
+ assert len(result) == 1, result
+ tcs = result[0].get("tool_calls") or []
+ assert len(tcs) == 1, result
+ assert tcs[0]["function"]["name"] == "code_execution"
+
+
+def test_native_part_code_execution_treated_as_server_side():
+ """Round 19: a Gemini `code_execution` card persists its replay
+ payload at `args.google.native_part` (no `_server_tool` marker on
+ pre-PR cards). The backend filter must still drop it for non-native
+ providers because it is a synthetic card, not a real user function."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ args_with_native_part = json.dumps(
+ {
+ "google": {
+ "native_part": {
+ "executableCode": {
+ "language": "PYTHON",
+ "code": "print(1)",
+ }
+ }
+ }
+ }
+ )
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_x",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": args_with_native_part,
+ },
+ }
+ ],
+ }
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "openai",
+ base_url = None,
+ )
+ assert result == [] or all(not (m.get("tool_calls") or []) for m in result), result
+
+
+def test_remote_image_fetch_attempt_cap_includes_failures(monkeypatch):
+ """Round 19: the per-request image fetch count cap must count
+ ATTEMPTS, not just successes. Otherwise a request with 100
+ failing/slow URLs runs 100 fetches each up to the 15s timeout."""
+ fetch_calls: list[str] = []
+
+ async def fake_fetch(url, fallback_mime, max_bytes = None):
+ fetch_calls.append(url)
+ return None
+
+ monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch)
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = _gemini_sse(
+ [
+ {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [{"text": "ok"}],
+ },
+ "finishReason": "STOP",
+ }
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 1,
+ "candidatesTokenCount": 1,
+ },
+ }
+ ]
+ ),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = _make_gemini_client()
+ image_parts = [
+ {
+ "type": "image_url",
+ "image_url": {"url": f"https://cdn.example.com/img{idx}.png"},
+ }
+ for idx in range(20)
+ ]
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "describe"},
+ *image_parts,
+ ],
+ }
+ ],
+ model = "gemini-2.5-flash",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 64,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ assert len(fetch_calls) <= 8, len(fetch_calls)
+
+
+def test_orphan_function_call_output_dropped_when_call_skipped(monkeypatch):
+ """Round 19: when a marked server-side builtin `function_call` is
+ dropped from OpenAI Responses input items, the matching role=tool
+ follow-up must also be dropped to avoid an orphan
+ `function_call_output`."""
+ captured: dict = {"input_items": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ body = json.loads(request.content.decode("utf-8"))
+ captured["input_items"] = body.get("input")
+ return httpx.Response(
+ 200,
+ content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openai",
+ base_url = "https://api.openai.com/v1",
+ api_key = "sk-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {"role": "user", "content": "search please"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {
+ "name": "web_search",
+ "arguments": json.dumps(
+ {"_server_tool": True, "query": "x"}
+ ),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "content": "result_text",
+ "tool_call_id": "call_b",
+ "name": "web_search",
+ },
+ {"role": "user", "content": "continue"},
+ ],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 1.0,
+ max_tokens = 16,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ items = captured["input_items"] or []
+ fn_calls = [i for i in items if i.get("type") == "function_call"]
+ fn_outs = [i for i in items if i.get("type") == "function_call_output"]
+ assert all(c.get("call_id") != "call_b" for c in fn_calls), items
+ assert all(o.get("call_id") != "call_b" for o in fn_outs), items
+
+
+def test_schema_multitype_union_with_null_preserves_anyof(monkeypatch):
+ """Round 19: a JSON Schema `"type": ["string","integer","null"]`
+ must be sanitized to anyOf:[{string},{integer}] + nullable:true.
+ Flattening to just `{"type":"string"}` silently drops the integer
+ branch and changes the function contract."""
+ captured = _capture_body(
+ monkeypatch,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "either": {"type": ["string", "integer", "null"]},
+ },
+ },
+ },
+ }
+ ],
+ )
+ decls = next(
+ t["functionDeclarations"]
+ for t in captured["body"].get("tools") or []
+ if "functionDeclarations" in t
+ )
+ either = decls[0]["parameters"]["properties"]["either"]
+ assert either.get("nullable") is True
+ inner = either.get("anyOf")
+ assert isinstance(inner, list) and len(inner) == 2, either
+ types = sorted(
+ b.get("type") for b in inner if isinstance(b, dict) and b.get("type")
+ )
+ assert types == ["integer", "string"], inner
+
+
+def test_invalid_gemini_model_rejected_before_image_fetch(monkeypatch):
+ """Round 19: invalid Gemini model IDs are rejected at the top of
+ `_stream_gemini`, BEFORE any user-controlled remote image fetch
+ runs."""
+ fetch_calls: list[str] = []
+
+ async def fake_fetch(url, fallback_mime, max_bytes = None):
+ fetch_calls.append(url)
+ return None
+
+ monkeypatch.setattr(ep_mod, "_safe_fetch_image_for_gemini", fake_fetch)
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = b"",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = _make_gemini_client()
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "hi"},
+ {
+ "type": "image_url",
+ "image_url": {"url": "https://cdn.example.com/x.png"},
+ },
+ ],
+ }
+ ],
+ model = "../cachedContents/leak",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 64,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ assert fetch_calls == [], fetch_calls
+
+
+def test_empty_assistant_turn_skipped_after_synthetic_tool_calls_dropped():
+ """Round 20: when `_filter_tool_calls` drops every synthetic
+ server-builtin tool_call on an empty-content assistant turn, the
+ whole message must be skipped. Forwarding
+ `{"role":"assistant","content":""}` is rejected by several
+ providers as an empty assistant turn."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ marked_args = json.dumps({"_server_tool": True, "kind": "image"})
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {
+ "name": "image_generation",
+ "arguments": marked_args,
+ },
+ }
+ ],
+ }
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+ for provider_type, base_url in [
+ ("openai", None),
+ ("gemini", "https://litellm.example/v1"),
+ ]:
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = provider_type,
+ base_url = base_url,
+ )
+ # The empty assistant turn (only a synthetic builtin) must
+ # NOT appear in the output at all.
+ assert result == [], (provider_type, result)
+
+
+def test_role_tool_dropped_when_matching_synthetic_call_filtered():
+ """Round 20: `_build_external_messages` drops the matching role=
+ tool follow-up when its tool_call was a synthetic builtin that
+ `_filter_tool_calls` removed. Otherwise the receiving provider
+ sees an orphan tool_result with no tool_call."""
+ from models.inference import ChatCompletionRequest
+ from routes.inference import _build_external_messages
+
+ marked_args = json.dumps({"_server_tool": True, "query": "x"})
+ payload = {
+ "model": "gpt-5.5",
+ "messages": [
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {
+ "name": "web_search",
+ "arguments": marked_args,
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "content": "result_text",
+ "tool_call_id": "call_b",
+ "name": "web_search",
+ },
+ {"role": "user", "content": "continue"},
+ ],
+ "stream": True,
+ }
+ req = ChatCompletionRequest.model_validate(payload)
+ result = _build_external_messages(
+ req.messages,
+ supports_vision = True,
+ provider_type = "openai",
+ base_url = None,
+ )
+ # Only the user "continue" message survives.
+ roles = [m.get("role") for m in result]
+ assert roles == ["user"], result
+
+
+def test_openrouter_no_synthetic_web_search_event_on_tool_choice_none(
+ monkeypatch,
+):
+ """Round 20: OpenRouter dispatcher must not emit synthetic
+ web_search tool_start / tool_end events when tool_choice="none";
+ otherwise the chat UI shows a search card for a search that
+ never happened."""
+ captured_events: list[dict] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = b"data: [DONE]\n\n",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openrouter",
+ base_url = "https://openrouter.ai/api/v1",
+ api_key = "sk-or-test",
+ )
+ async for line in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "openai/gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search"],
+ tool_choice = "none",
+ ):
+ if not line.startswith("data: "):
+ continue
+ payload = line[len("data: ") :].strip()
+ if not payload or payload == "[DONE]":
+ continue
+ try:
+ obj = json.loads(payload)
+ except Exception:
+ continue
+ # Backend emits synthetic tool events as a top-level
+ # `_toolEvent` on the SSE payload (not nested inside
+ # `delta`). Read both shapes so a future format change
+ # cannot mask this regression.
+ evt = obj.get("_toolEvent")
+ if isinstance(evt, dict):
+ captured_events.append(evt)
+ for ch in obj.get("choices") or []:
+ delta = ch.get("delta") or {}
+ nested = delta.get("_toolEvent") if isinstance(delta, dict) else None
+ if isinstance(nested, dict):
+ captured_events.append(nested)
+ await client.close()
+
+ _drive(run())
+ # No synthetic web_search tool_start / tool_end emitted.
+ assert all(
+ e.get("tool_name") != "web_search" for e in captured_events
+ ), captured_events
+
+
+def test_anthropic_role_tool_list_content_translates_to_tool_result(
+ monkeypatch,
+):
+ """Round 20: an OpenAI-shape role=tool message with list content
+ (`content=[{"type":"text","text":"result"}]`) must be translated
+ into Anthropic's native tool_result block, not forwarded as an
+ invalid role=tool message."""
+ captured: dict = {"messages": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ body = json.loads(request.content.decode("utf-8"))
+ captured["messages"] = body.get("messages")
+ return httpx.Response(
+ 200,
+ content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "anthropic",
+ base_url = "https://api.anthropic.com",
+ api_key = "sk-ant-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {"role": "user", "content": "look up X"},
+ {
+ "role": "assistant",
+ "content": "let me check",
+ "tool_calls": [
+ {
+ "id": "call_a",
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "arguments": '{"q":"x"}',
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "content": [{"type": "text", "text": "result_text"}],
+ "tool_call_id": "call_a",
+ "name": "lookup",
+ },
+ {"role": "user", "content": "summarise"},
+ ],
+ model = "claude-sonnet-4-5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 64,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ msgs = captured["messages"] or []
+ assert all(m.get("role") != "tool" for m in msgs), msgs
+ tool_results: list[dict] = []
+ for m in msgs:
+ if m.get("role") == "user" and isinstance(m.get("content"), list):
+ tool_results.extend(
+ b for b in m["content"] if b.get("type") == "tool_result"
+ )
+ assert any(
+ tr.get("tool_use_id") == "call_a" and tr.get("content") == "result_text"
+ for tr in tool_results
+ ), msgs
+
+
+def test_data_url_non_image_mime_dropped(monkeypatch):
+ """Round 20: a `data:text/html;base64,...` image_url must be
+ dropped from the outbound Gemini body, not forwarded as
+ `inlineData.mimeType="text/html"` which Gemini rejects."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "look"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "data:text/html;base64,PGgxPmhpPC9oMT4=",
+ },
+ },
+ ],
+ }
+ ],
+ )
+ parts = captured["body"]["contents"][-1]["parts"]
+ assert not any("inlineData" in p for p in parts), parts
+
+
+def test_youtube_filedata_uses_video_mime(monkeypatch):
+ """Round 20: YouTube `fileData.fileUri` must declare a video
+ mimeType, not `image/jpeg` guessed from the URL path."""
+ captured = _capture_body(
+ monkeypatch,
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "summarise"},
+ {
+ "type": "image_url",
+ "image_url": {
+ "url": "https://www.youtube.com/watch?v=abc",
+ },
+ },
+ ],
+ }
+ ],
+ )
+ parts = captured["body"]["contents"][-1]["parts"]
+ yt = next((p for p in parts if "fileData" in p), None)
+ assert yt is not None, parts
+ assert yt["fileData"]["mimeType"].startswith("video/"), yt
+
+
+def test_openai_responses_assistant_text_serialized_before_function_call(
+ monkeypatch,
+):
+ """Round 20: in OpenAI Responses history, the assistant's
+ visible text for a turn that ALSO emitted a function_call must
+ serialize BEFORE the function_call item, matching the prior
+ response.output sequence. Otherwise function_call_output (the
+ role=tool follow-up) appears to follow an unrelated assistant
+ message."""
+ captured: dict = {"input_items": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ body = json.loads(request.content.decode("utf-8"))
+ captured["input_items"] = body.get("input")
+ return httpx.Response(
+ 200,
+ content = b'data: {"type":"response.completed","response":{"output":[],"usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openai",
+ base_url = "https://api.openai.com/v1",
+ api_key = "sk-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [
+ {"role": "user", "content": "weather?"},
+ {
+ "role": "assistant",
+ "content": "Let me check that.",
+ "tool_calls": [
+ {
+ "id": "call_w",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": "{}",
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "content": "sunny",
+ "tool_call_id": "call_w",
+ "name": "get_weather",
+ },
+ {"role": "user", "content": "thanks"},
+ ],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 1.0,
+ max_tokens = 16,
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+
+ items = captured["input_items"] or []
+ types = [i.get("type") or i.get("role") for i in items]
+ # Expected order:
+ # user ("weather?")
+ # assistant ("Let me check that.")
+ # function_call (get_weather)
+ # function_call_output (sunny)
+ # user ("thanks")
+ assert types == [
+ "user",
+ "assistant",
+ "function_call",
+ "function_call_output",
+ "user",
+ ], items
+
+
+def test_gemini_tool_choice_none_disables_image_generation(monkeypatch):
+ """Round 21: `tool_choice="none"` must also flip the implicit
+ image-generation hosted tool off on image-tier models. Otherwise
+ `responseModalities=["TEXT","IMAGE"]` still rides on the outbound
+ body and the provider can generate (and bill for) image output
+ despite the explicit OpenAI tool opt-out."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation"],
+ tool_choice = "none",
+ )
+ body = captured["body"]
+ assert body["generationConfig"].get("responseModalities") == ["TEXT"], body
+
+
+def test_gemini_forced_function_tool_choice_drops_hosted_builtins(monkeypatch):
+ """Round 21: forced-function `tool_choice` (e.g.
+ `{"type":"function","function":{"name":"lookup"}}`) must suppress
+ hosted Google Search / code execution. Gemini's toolConfig only
+ constrains function declarations, not hosted tools, so leaving
+ `googleSearch`/`codeExecution` in `tools[]` lets them fire despite
+ the caller pinning a specific user function."""
+ captured = _capture_body(
+ monkeypatch,
+ enabled_tools = ["web_search", "code_execution"],
+ tools = [
+ {
+ "type": "function",
+ "function": {"name": "lookup", "parameters": {"type": "object"}},
+ }
+ ],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup"},
+ },
+ )
+ body = captured["body"]
+ tool_kinds = [list(t.keys())[0] for t in (body.get("tools") or [])]
+ assert "googleSearch" not in tool_kinds, body
+ assert "codeExecution" not in tool_kinds, body
+ # User function declaration still survives.
+ assert "functionDeclarations" in tool_kinds, body
+
+
+def test_gemini_forced_function_tool_choice_drops_image_generation(monkeypatch):
+ """Round 21: forced-function `tool_choice` must also flip the
+ implicit image-generation hosted tool off on image-tier models."""
+ captured = _capture_body(
+ monkeypatch,
+ model = "gemini-2.5-flash-image",
+ enabled_tools = ["image_generation"],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup"},
+ },
+ tools = [
+ {
+ "type": "function",
+ "function": {"name": "lookup", "parameters": {"type": "object"}},
+ }
+ ],
+ )
+ body = captured["body"]
+ assert body["generationConfig"].get("responseModalities") == ["TEXT"], body
+
+
+def test_gemini_code_execution_native_part_list_replays_per_part_signatures(
+ monkeypatch,
+):
+ """Round 21: merged code-execution history must replay per-part
+ `thoughtSignature`s, not fan one top-level signature across every
+ native subpart. Gemini 3 strict validators reject a signature
+ placed on the wrong part."""
+ history = [
+ {"role": "user", "content": "plot 1+1"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "call_a",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": "{}",
+ },
+ "extra_content": {
+ "google": {
+ "native_part": {
+ "parts": [
+ {
+ "executableCode": {
+ "id": "code_a",
+ "language": "PYTHON",
+ "code": "print(1+1)",
+ },
+ "thoughtSignature": "SIG-EXEC",
+ },
+ {
+ "codeExecutionResult": {
+ "id": "res_a",
+ "outcome": "OUTCOME_OK",
+ "output": "2\n",
+ },
+ },
+ ],
+ },
+ },
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_a",
+ "name": "code_execution",
+ "content": "2",
+ },
+ {"role": "user", "content": "next"},
+ ]
+ captured = _capture_body(monkeypatch, messages = history)
+ contents = captured["body"]["contents"]
+ # Locate the assistant turn replayed as native code-exec parts.
+ assistant_turn = next(c for c in contents if c["role"] == "model")
+ parts = assistant_turn["parts"]
+ exec_parts = [p for p in parts if "executableCode" in p]
+ result_parts = [p for p in parts if "codeExecutionResult" in p]
+ assert exec_parts and result_parts, parts
+ assert exec_parts[0].get("thoughtSignature") == "SIG-EXEC", exec_parts[0]
+ # codeExecutionResult had no signature -- must NOT inherit one.
+ assert "thoughtSignature" not in result_parts[0], result_parts[0]
+
+
+def test_gemini_code_execution_legacy_merged_signature_only_on_executable(
+ monkeypatch,
+):
+ """Round 21: backward compatibility for pre-round-21 persisted
+ history that stored merged `native_part` as a single object plus a
+ top-level `thoughtSignature`. The replay branch must attach that
+ signature only to `executableCode` (where Gemini 3 emits it), not
+ fan it across `codeExecutionResult` / `inlineData`."""
+ history = [
+ {"role": "user", "content": "plot 1+1"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "call_b",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": "{}",
+ },
+ "extra_content": {
+ "google": {
+ "native_part": {
+ "executableCode": {
+ "id": "code_b",
+ "language": "PYTHON",
+ "code": "print(1+1)",
+ },
+ "codeExecutionResult": {
+ "id": "res_b",
+ "outcome": "OUTCOME_OK",
+ "output": "2\n",
+ },
+ "thoughtSignature": "LEGACY-SIG",
+ },
+ },
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_b",
+ "name": "code_execution",
+ "content": "2",
+ },
+ {"role": "user", "content": "next"},
+ ]
+ captured = _capture_body(monkeypatch, messages = history)
+ contents = captured["body"]["contents"]
+ assistant_turn = next(c for c in contents if c["role"] == "model")
+ exec_parts = [p for p in assistant_turn["parts"] if "executableCode" in p]
+ result_parts = [p for p in assistant_turn["parts"] if "codeExecutionResult" in p]
+ assert exec_parts[0].get("thoughtSignature") == "LEGACY-SIG", exec_parts[0]
+ assert "thoughtSignature" not in result_parts[0], result_parts[0]
+
+
+def test_gemini_role_tool_list_content_flattens_to_result_text(monkeypatch):
+ """Round 21: OpenAI-shape role=tool messages may carry list content
+ like `[{"type":"text","text":"result"}]`. Forwarding those parts
+ verbatim into `functionResponse.response.result` yields a list of
+ content-part objects instead of the actual tool output text."""
+ history = [
+ {"role": "user", "content": "look up"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_1",
+ "type": "function",
+ "function": {
+ "name": "lookup",
+ "arguments": json.dumps({"q": "x"}),
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_1",
+ "name": "lookup",
+ "content": [{"type": "text", "text": "answer-text"}],
+ },
+ {"role": "user", "content": "next"},
+ ]
+ captured = _capture_body(monkeypatch, messages = history)
+ contents = captured["body"]["contents"]
+ fn_response = None
+ for c in contents:
+ for p in c.get("parts") or []:
+ if isinstance(p, dict) and "functionResponse" in p:
+ fn_response = p["functionResponse"]
+ break
+ if fn_response:
+ break
+ assert fn_response is not None, contents
+ assert fn_response["response"] == {"result": "answer-text"}, fn_response
+
+
+def test_safe_fetch_image_threads_per_request_byte_budget(monkeypatch):
+ """Round 21: the aggregate per-request byte cap must be passed into
+ `_safe_fetch_image_for_gemini` so an oversize URL is refused via
+ Content-Length (short-circuit) rather than fully downloaded then
+ discarded after the fact."""
+ import socket
+
+ captured: dict = {"reads": 0, "content_length_seen": None}
+
+ original_getaddrinfo = socket.getaddrinfo
+
+ def fake_getaddrinfo(host, *args, **kwargs):
+ if host == "cdn.example.com":
+ return [
+ (
+ socket.AF_INET,
+ socket.SOCK_STREAM,
+ 0,
+ "",
+ ("8.8.8.8", 0),
+ )
+ ]
+ return original_getaddrinfo(host, *args, **kwargs)
+
+ monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo)
+
+ class _StubResp:
+ status = 200
+ # Declared 5 MiB, but caller passes a 1 MiB remaining budget.
+ headers = {
+ "content-type": "image/png",
+ "content-length": str(5 * 1024 * 1024),
+ }
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *a):
+ return False
+
+ def read(self, _n = None):
+ captured["reads"] += 1
+ return b"\x00" * (5 * 1024 * 1024)
+
+ class _StubOpener:
+ def open(self, req, timeout = None):
+ return _StubResp()
+
+ monkeypatch.setattr(
+ "urllib.request.build_opener", lambda *_args, **_kw: _StubOpener()
+ )
+
+ res = _drive(
+ ep_mod._safe_fetch_image_for_gemini(
+ "https://cdn.example.com/big.png",
+ "image/png",
+ max_bytes = 1 * 1024 * 1024,
+ )
+ )
+ assert res is None
+ # Refused via Content-Length pre-check, never read.
+ assert captured["reads"] == 0
+
+
+def test_openai_chat_delta_type_includes_tool_calls_and_extra_content():
+ """Round 21: the frontend `OpenAIChatDelta` interface must expose
+ `tool_calls` and `extra_content` so TypeScript callers can consume
+ the Gemini-native stream fields without `any` casts. This test is
+ a static-string assertion against the .ts source; mirrors how other
+ frontend wire-contract tests are pinned from the backend suite."""
+ import os
+
+ here = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ types_path = os.path.join(
+ here, "frontend", "src", "features", "chat", "types", "api.ts"
+ )
+ with open(types_path, "r", encoding = "utf-8") as f:
+ src = f.read()
+ assert "tool_calls?: OpenAIToolCallPart[]" in src, src[:200]
+ assert "extra_content?: Record" in src, src[:200]
+ assert "boolean | string | null" in src, src[:200]
+
+
+def test_anthropic_forced_function_tool_choice_drops_hosted_tools(monkeypatch):
+ """Round 22: forced-function tool_choice must suppress Anthropic
+ hosted builtins the same way it does for Gemini. Pinning a user
+ function (`tool_choice={"type":"function","function":{"name":...}}`)
+ while also passing `enabled_tools=["web_search","web_fetch",
+ "code_execution"]` should not still fire those server-side."""
+ captured: dict = {"body": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "anthropic",
+ base_url = "https://api.anthropic.com",
+ api_key = "sk-ant-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "claude-sonnet-4-5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search", "web_fetch", "code_execution"],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup_record"},
+ },
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ body = captured["body"] or {}
+ # No hosted tools should be in the body — only the caller's user-
+ # function declarations (which this test doesn't pass any of).
+ tools = body.get("tools") or []
+ hosted_tool_names = {"web_search", "web_fetch", "code_execution"}
+ for tool in tools:
+ assert tool.get("name") not in hosted_tool_names, body
+
+
+def test_openrouter_forced_function_tool_choice_drops_web_plugin(monkeypatch):
+ """Round 22: forced-function tool_choice must drop the OpenRouter
+ web plugin too — caller pinned a user function, OpenRouter must not
+ still attach the hosted web-search plugin."""
+ captured: dict = {"body": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = b"data: [DONE]\n\n",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openrouter",
+ base_url = "https://openrouter.ai/api/v1",
+ api_key = "sk-or-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "openai/gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search"],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup_record"},
+ },
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ body = captured["body"] or {}
+ assert body.get("plugins") in (None, []), body
+
+
+def test_kimi_forced_function_tool_choice_skips_web_search_helper(monkeypatch):
+ """Round 22: forced-function tool_choice plus enabled_tools=
+ ["web_search"] on Kimi must NOT route into `_stream_kimi_web_search`.
+ Caller pinned a user function; hosted $web_search should be
+ suppressed for the same privacy/billing reason."""
+ routed_to_helper = {"called": False}
+
+ async def fake_helper(self, *args, **kwargs): # noqa: ARG001
+ routed_to_helper["called"] = True
+ if False:
+ yield "" # pragma: no cover
+
+ monkeypatch.setattr(
+ ExternalProviderClient,
+ "_stream_kimi_web_search",
+ fake_helper,
+ )
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = b"data: [DONE]\n\n",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "kimi",
+ base_url = "https://api.moonshot.ai/v1",
+ api_key = "sk-kimi-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "kimi-k2.6",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search"],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup_record"},
+ },
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ assert not routed_to_helper["called"]
+
+
+def test_openai_responses_forced_function_tool_choice_drops_hosted_tools(monkeypatch):
+ """Round 23: forced-function tool_choice on the OpenAI Responses
+ path must suppress hosted builtins (web_search, shell,
+ image_generation) the same way it does for Gemini / Anthropic /
+ OpenRouter / Kimi. User-defined function tools still flow through
+ so the pinned function can resolve."""
+ captured: dict = {"body": None}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = b"event: response.completed\ndata: {}\n\n",
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openai",
+ base_url = "https://api.openai.com/v1",
+ api_key = "sk-openai-test",
+ )
+ async for _ in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "gpt-5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search", "code_execution", "image_generation"],
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup_record",
+ "parameters": {"type": "object", "properties": {}},
+ },
+ },
+ ],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup_record"},
+ },
+ ):
+ pass
+ await client.close()
+
+ _drive(run())
+ body = captured["body"] or {}
+ tools = body.get("tools") or []
+ hosted_types = {"web_search", "shell", "image_generation"}
+ hosted_seen = {t.get("type") for t in tools if isinstance(t, dict)}
+ assert not (hosted_seen & hosted_types), body
+ # The user function declaration must still be present so the pin
+ # has something to target.
+ user_function_seen = any(
+ isinstance(t, dict) and t.get("type") == "function" for t in tools
+ )
+ assert user_function_seen, body
+ # And the forced-function tool_choice must be forwarded in Responses
+ # shape: `{type:"function", name:"..."}`.
+ tc = body.get("tool_choice")
+ assert isinstance(tc, dict) and tc.get("type") == "function", body
+ assert tc.get("name") == "lookup_record", body
+
+
+def test_strip_provider_synthetic_tool_history_drops_text_only_extra_content():
+ """Round 24: a plain text Gemini reply (no tool_calls) carrying
+ `extra_content.google.thought_signature` must still have that
+ metadata stripped before being forwarded to a local llama-server
+ backend. Without this, switching a Gemini thread mid-stream to a
+ local GGUF model leaks Gemini-only fields to llama-server."""
+ from routes.inference import _strip_provider_synthetic_tool_history
+
+ messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": "Hello!",
+ "extra_content": {"google": {"thought_signature": "SIG_ABC"}},
+ },
+ {"role": "user", "content": "now in pirate voice"},
+ ]
+ out = _strip_provider_synthetic_tool_history(messages)
+ # Same three turns, but the assistant's `extra_content` is gone.
+ assert [m["role"] for m in out] == ["user", "assistant", "user"]
+ assistant = out[1]
+ assert "extra_content" not in assistant, assistant
+ assert assistant["content"] == "Hello!"
+
+
+def test_validate_and_resolve_host_blocks_shared_address_space():
+ """Round 24 SSRF P1: 100.64.0.0/10 carrier-grade NAT addresses are
+ `is_private=False` AND `is_global=False` per Python's ipaddress
+ docs. The previous denylist (is_private/loopback/link_local/etc.)
+ missed them. Adding `not ip.is_global` as the primary gate covers
+ all non-public ranges, current and future."""
+ import socket as _socket
+ from core.inference import tools as _tools
+
+ orig_getaddrinfo = _socket.getaddrinfo
+
+ def fake_getaddrinfo(hostname, port, *args, **kwargs):
+ if hostname == "shared.example":
+ return [
+ (
+ _socket.AF_INET,
+ _socket.SOCK_STREAM,
+ 0,
+ "",
+ ("100.64.0.1", port),
+ ),
+ ]
+ return orig_getaddrinfo(hostname, port, *args, **kwargs)
+
+ _socket.getaddrinfo = fake_getaddrinfo
+ try:
+ ok, reason, _ip = _tools._validate_and_resolve_host("shared.example", 443)
+ finally:
+ _socket.getaddrinfo = orig_getaddrinfo
+ assert ok is False, (ok, reason)
+ assert "non-public" in reason.lower() or "100.64.0.1" in reason
+
+
+def test_gemini_custom_oai_compat_base_skips_native_allowlist():
+ """Round 24: a custom Gemini OAI-compatible base (LiteLLM/proxy)
+ must NOT have its model list filtered through the native Gemini
+ allowlist regex. A LiteLLM gateway returning
+ `["google/gemini-2.5-flash", "my-team/gemini", "gemini-2.5-flash"]`
+ should be passed through; the native filter would strip the
+ prefixed IDs even though the chat dispatch routes them via the
+ OpenAI-compatible client."""
+ import asyncio as _asyncio
+
+ from routes import providers as _providers
+ from routes.providers import (
+ ProviderModelsRequest,
+ list_provider_models,
+ )
+
+ captured: dict = {"base": None}
+
+ class _FakeClient:
+ def __init__(self, *, base_url, **kwargs):
+ captured["base"] = base_url
+
+ async def list_models(self):
+ return [
+ {"id": "google/gemini-2.5-flash"},
+ {"id": "my-team/gemini"},
+ {"id": "gemini-2.5-flash"},
+ ]
+
+ async def close(self):
+ return None
+
+ orig = _providers.ExternalProviderClient
+ _providers.ExternalProviderClient = _FakeClient
+ try:
+ req = ProviderModelsRequest(
+ provider_type = "gemini",
+ base_url = "https://litellm.example/v1",
+ )
+ result = _asyncio.run(list_provider_models(req, current_subject = "unsloth"))
+ finally:
+ _providers.ExternalProviderClient = orig
+ ids = {m.id for m in result}
+ # All three IDs survive — the native allowlist was bypassed.
+ assert "google/gemini-2.5-flash" in ids, ids
+ assert "my-team/gemini" in ids, ids
+ assert "gemini-2.5-flash" in ids, ids
+
+
+def test_strip_provider_synthetic_tool_history_drops_synthetic_only():
+ """Round 22: switching a thread from native Gemini (code_execution
+ / image_generation tool_cards in history) to a local GGUF backend
+ must strip the synthetic tool_calls + matching role=tool replies
+ before llama-server sees them. Real user-function tool_calls and
+ their matching tool replies must survive."""
+ from routes.inference import _strip_provider_synthetic_tool_history
+
+ messages = [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": "let me run it",
+ "tool_calls": [
+ {
+ "id": "synth_ce_1",
+ "type": "function",
+ "function": {
+ "name": "code_execution",
+ "arguments": json.dumps(
+ {
+ "_server_tool": True,
+ "google": {"native_part": {"parts": []}},
+ }
+ ),
+ },
+ "extra_content": {"google": {"thought_signature": "abc"}},
+ },
+ {
+ "id": "real_lookup",
+ "type": "function",
+ "function": {
+ "name": "lookup_user",
+ "arguments": json.dumps({"id": 42}),
+ },
+ },
+ ],
+ "extra_content": {"google": {"thought_signature": "msglevel"}},
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "synth_ce_1",
+ "content": "Gemini-only result text",
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "real_lookup",
+ "content": '{"name": "alice"}',
+ },
+ ]
+ out = _strip_provider_synthetic_tool_history(messages)
+ assistant = next(m for m in out if m.get("role") == "assistant")
+ tcs = assistant["tool_calls"]
+ assert len(tcs) == 1, tcs
+ assert tcs[0]["id"] == "real_lookup"
+ assert "extra_content" not in tcs[0]
+ assert "extra_content" not in assistant
+ tool_msgs = [m for m in out if m.get("role") == "tool"]
+ assert len(tool_msgs) == 1
+ assert tool_msgs[0]["tool_call_id"] == "real_lookup"
+
+
+def test_strip_provider_synthetic_tool_history_drops_empty_assistant():
+ """If every tool_call was synthetic and the assistant turn had no
+ content, the entire turn must be dropped (llama-server rejects
+ empty assistant messages with no tool_calls)."""
+ from routes.inference import _strip_provider_synthetic_tool_history
+
+ messages = [
+ {"role": "user", "content": "draw a sloth"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "synth_imggen",
+ "type": "function",
+ "function": {
+ "name": "image_generation",
+ "arguments": json.dumps(
+ {
+ "google": {
+ "native_part": {
+ "parts": [
+ {
+ "inlineData": {
+ "mimeType": "image/png",
+ "data": "Zm9v",
+ }
+ }
+ ]
+ }
+ }
+ }
+ ),
+ },
+ }
+ ],
+ },
+ {"role": "tool", "tool_call_id": "synth_imggen", "content": "(image)"},
+ {"role": "user", "content": "now try in pirate voice"},
+ ]
+ out = _strip_provider_synthetic_tool_history(messages)
+ roles = [m.get("role") for m in out]
+ # The synthetic assistant + its tool reply are both gone; only the
+ # two user turns survive.
+ assert roles == ["user", "user"], out
+
+
+def test_openrouter_no_synthetic_web_search_event_on_forced_function_tool_choice(
+ monkeypatch,
+):
+ """Round 22 sibling of the round-20 `tool_choice='none'` test: when
+ the caller forces a specific function via `tool_choice={"type":
+ "function", ...}` AND passes `enabled_tools=["web_search"]`, the
+ OpenRouter path must NOT synthesize a fake `web_search` tool card.
+ The plugin was not attached upstream so the UI must not see a
+ server-tool card."""
+ captured_events: list[dict] = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ content = (
+ b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\n'
+ b"data: [DONE]\n\n"
+ ),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http(monkeypatch, handler)
+
+ async def run():
+ client = ExternalProviderClient(
+ provider_type = "openrouter",
+ base_url = "https://openrouter.ai/api/v1",
+ api_key = "sk-or-test",
+ )
+ async for line in client.stream_chat_completion(
+ messages = [{"role": "user", "content": "hi"}],
+ model = "openai/gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = 16,
+ enabled_tools = ["web_search"],
+ tool_choice = {
+ "type": "function",
+ "function": {"name": "lookup_record"},
+ },
+ ):
+ payload = line.strip().removeprefix("data: ")
+ if payload and payload != "[DONE]":
+ try:
+ captured_events.append(json.loads(payload))
+ except Exception:
+ pass
+ await client.close()
+
+ _drive(run())
+ for evt in captured_events:
+ for choice in evt.get("choices") or []:
+ delta = choice.get("delta") or {}
+ extra = delta.get("extra_content") or {}
+ tool_event = extra.get("toolEvent") if isinstance(extra, dict) else None
+ if isinstance(tool_event, dict):
+ assert tool_event.get("tool_name") != "web_search", evt
diff --git a/studio/backend/tests/test_gguf_completion_usage.py b/studio/backend/tests/test_gguf_completion_usage.py
new file mode 100644
index 0000000000..b8cfaee7c7
--- /dev/null
+++ b/studio/backend/tests/test_gguf_completion_usage.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+"""Regression tests for GGUF non-streaming chat completion usage."""
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from auth.authentication import get_current_subject
+import routes.inference as inference_route
+
+
+class _GgufBackend:
+ is_loaded = True
+ model_identifier = "test/model.gguf"
+ _is_audio = False
+ is_vision = False
+ supports_tools = False
+
+ def __init__(self, usage):
+ self.usage = usage
+
+ def generate_chat_completion(self, **kwargs):
+ yield "answer"
+ yield {
+ "type": "metadata",
+ "usage": self.usage,
+ "timings": {"prompt_n": 23, "predicted_n": 1283},
+ }
+
+
+def _request_completion(monkeypatch, usage):
+ monkeypatch.setattr(
+ inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage)
+ )
+ monkeypatch.setattr(
+ inference_route, "_effective_enable_tools", lambda payload: False
+ )
+
+ app = FastAPI()
+ app.include_router(inference_route.router)
+ app.dependency_overrides[get_current_subject] = lambda: "test-user"
+
+ return TestClient(app).post(
+ "/chat/completions",
+ json = {
+ "messages": [{"role": "user", "content": "Why is the sky blue?"}],
+ "stream": False,
+ },
+ )
+
+
+def test_non_streaming_gguf_completion_includes_generated_usage(monkeypatch):
+ response = _request_completion(
+ monkeypatch,
+ {"prompt_tokens": 23, "completion_tokens": 1283, "total_tokens": 1306},
+ )
+
+ assert response.status_code == 200
+ assert response.json()["usage"] == {
+ "prompt_tokens": 23,
+ "completion_tokens": 1283,
+ "total_tokens": 1306,
+ }
+
+
+def test_non_streaming_gguf_completion_defaults_nullable_usage_to_zero(monkeypatch):
+ response = _request_completion(
+ monkeypatch,
+ {"prompt_tokens": None, "completion_tokens": 1283, "total_tokens": None},
+ )
+
+ assert response.status_code == 200
+ assert response.json()["usage"] == {
+ "prompt_tokens": 0,
+ "completion_tokens": 1283,
+ "total_tokens": 0,
+ }
diff --git a/studio/backend/tests/test_gguf_routing.py b/studio/backend/tests/test_gguf_routing.py
new file mode 100644
index 0000000000..1b299cba19
--- /dev/null
+++ b/studio/backend/tests/test_gguf_routing.py
@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+Tests for GGUF routing in detect_gguf_model.
+
+Regression test for the bug where a .gguf file temporarily appears
+inaccessible on Windows during llama-server process teardown, causing
+is_file() to return False and the model to be routed to the transformers
+backend instead of llama-server.
+"""
+
+import sys
+import os
+import types
+from pathlib import Path
+from unittest.mock import patch
+
+# Stub structlog before importing backend modules (mirrors other tests in this suite)
+if "structlog" not in sys.modules:
+
+ class _DummyLogger:
+ def __getattr__(self, _):
+ return lambda *a, **k: None
+
+ sys.modules["structlog"] = types.SimpleNamespace(
+ get_logger = lambda *a, **k: _DummyLogger(),
+ BoundLogger = _DummyLogger,
+ )
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from utils.models.model_config import detect_gguf_model
+
+
+def test_detects_gguf_file_normally(tmp_path):
+ """Normal case: .gguf file exists and is accessible."""
+ gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf"
+ gguf.write_bytes(b"")
+ result = detect_gguf_model(str(gguf))
+ assert result is not None
+ assert result.endswith("gpt-oss-20b-MXFP4.gguf")
+
+
+def test_detects_gguf_when_stat_raises_oserror(tmp_path):
+ """
+ Regression: on Windows, both is_file() and exists() call stat() internally.
+ During the brief lock window after llama-server is killed, stat() raises
+ OSError, causing both to return False. detect_gguf_model must still route
+ to llama-server based on the file extension alone.
+ """
+ gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf"
+ gguf.write_bytes(b"")
+
+ original_stat = Path.stat
+
+ def flaky_stat(self, **kwargs):
+ if self.suffix.lower() == ".gguf":
+ raise OSError("file temporarily inaccessible (Windows lock window)")
+ return original_stat(self, **kwargs)
+
+ with patch.object(Path, "stat", flaky_stat):
+ result = detect_gguf_model(str(gguf))
+
+ assert result is not None, (
+ "detect_gguf_model returned None when stat() raised OSError. "
+ "This causes the model to fall through to the transformers backend."
+ )
+
+
+def test_does_not_detect_mmproj_as_main_model(tmp_path):
+ """mmproj files must never be returned as the primary model."""
+ mmproj = tmp_path / "mmproj-model-f16.gguf"
+ mmproj.write_bytes(b"")
+ result = detect_gguf_model(str(mmproj))
+ assert result is None
+
+
+def test_detects_gguf_in_directory(tmp_path):
+ """Directory containing a .gguf file is resolved to that file."""
+ gguf = tmp_path / "model-Q4_K_M.gguf"
+ gguf.write_bytes(b"")
+ result = detect_gguf_model(str(tmp_path))
+ assert result is not None
+ assert result.endswith("model-Q4_K_M.gguf")
+
+
+def test_directory_named_like_gguf_scans_inside(tmp_path):
+ """A directory named *.gguf resolves the real .gguf inside, not itself."""
+ gguf_dir = tmp_path / "mymodel.gguf"
+ gguf_dir.mkdir()
+ inner = gguf_dir / "model-Q4_K_M.gguf"
+ inner.write_bytes(b"")
+ result = detect_gguf_model(str(gguf_dir))
+ assert result is not None
+ assert result.endswith("model-Q4_K_M.gguf")
+
+
+def test_returns_none_for_non_gguf_path(tmp_path):
+ """Non-.gguf paths with no .gguf files inside return None."""
+ result = detect_gguf_model(str(tmp_path))
+ assert result is None
diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
new file mode 100644
index 0000000000..5d2d672890
--- /dev/null
+++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
@@ -0,0 +1,430 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets.
+
+Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade
+GitHub API are stubbed out so the suite runs without internet access and is
+not subject to rate limits.
+"""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+_studio = Path(__file__).resolve().parent.parent.parent
+if str(_studio) not in sys.path:
+ sys.path.insert(0, str(_studio))
+
+_mod = importlib.import_module("install_llama_prebuilt")
+HostInfo = _mod.HostInfo
+resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None)
+_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None)
+
+if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
+ pytest.skip("PR symbols not present - check branch", allow_module_level = True)
+
+
+@pytest.fixture(autouse = True)
+def _clear_lemonade_release_cache():
+ """Prevent cross-test pollution of the lemonade release lru_cache when
+ future tests vary the fetch_json mock return value."""
+ _cache = getattr(_mod, "_fetch_lemonade_release_cached", None)
+ if _cache is not None and hasattr(_cache, "cache_clear"):
+ _cache.cache_clear()
+ yield
+ if _cache is not None and hasattr(_cache, "cache_clear"):
+ _cache.cache_clear()
+
+
+_STUB_TAG = "b1262"
+_STUB_OS_PREFIXES = ("ubuntu", "windows")
+_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X")
+
+
+def _stub_lemonade_release() -> dict:
+ """Minimal lemonade release payload covering all supported GPU/OS combinations."""
+ assets = [
+ {
+ "name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip",
+ "browser_download_url": (
+ f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/"
+ f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip"
+ ),
+ }
+ for prefix in _STUB_OS_PREFIXES
+ for family in _STUB_FAMILIES
+ ]
+ return {"tag_name": _STUB_TAG, "assets": assets}
+
+
+def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo:
+ return HostInfo(
+ system = "Windows" if windows else "Linux",
+ machine = "amd64" if windows else "x86_64",
+ is_windows = windows,
+ is_linux = not windows,
+ 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,
+ has_rocm = True,
+ rocm_gfx_target = gfx_target,
+ )
+
+
+def _lookup_family(gfx: str) -> str | None:
+ for prefix, family in _LEMONADE_GFX_FAMILIES:
+ if gfx.startswith(prefix):
+ return family
+ return None
+
+
+# ---------------------------------------------------------------------------
+# GPU family mapping
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "gfx,expected_family",
+ [
+ ("gfx1151", "gfx1151"),
+ ("gfx1150", "gfx1150"),
+ ("gfx1201", "gfx120X"),
+ ("gfx1200", "gfx120X"),
+ ("gfx1100", "gfx110X"),
+ ("gfx1030", "gfx103X"),
+ ],
+)
+def test_gpu_family_mapping(gfx, expected_family):
+ assert _lookup_family(gfx) == expected_family
+
+
+def test_unknown_gpu_not_in_families():
+ assert _lookup_family("gfx999") is None
+
+
+# ---------------------------------------------------------------------------
+# Asset resolution - hits real lemonade GitHub API
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "gfx,os_prefix,windows",
+ [
+ ("gfx1151", "ubuntu", False),
+ ("gfx1150", "ubuntu", False),
+ ("gfx1201", "ubuntu", False),
+ ("gfx1100", "ubuntu", False),
+ ("gfx1030", "ubuntu", False),
+ ("gfx1151", "windows", True),
+ ("gfx1100", "windows", True),
+ ],
+)
+def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows):
+ host = _make_rocm_host(gfx, windows = windows)
+ with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
+ result = resolve_lemonade_rocm_choice(
+ host, os_prefix, "default", llama_tag = "latest"
+ )
+ assert (
+ result is not None
+ ), f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})"
+ assert _lookup_family(gfx) in result.name
+ assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm")
+
+
+def test_unknown_gpu_falls_through_to_upstream():
+ host = _make_rocm_host("gfx999")
+ result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest")
+ assert result is None
+
+
+# ---------------------------------------------------------------------------
+# Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts.
+# This is the path setup.sh actually invokes (via --simple-policy), so the
+# lemonade integration is useless if it isn't wired in here.
+# ---------------------------------------------------------------------------
+
+direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None)
+direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None)
+
+
+def _stub_unsloth_release(release_tag: str = "b9022") -> dict:
+ # Minimal payload that parse_direct_linux_release_bundle accepts. It
+ # requires at least one `app-{label}-linux-x64*.tar.gz` asset for the
+ # bundle to be recognised; we ship a bare CPU one so the planner has a
+ # baseline non-ROCm attempt to fall through to.
+ asset_name = f"app-{release_tag}-linux-x64.tar.gz"
+ return {
+ "tag_name": release_tag,
+ "name": release_tag,
+ "assets": [
+ {
+ "name": asset_name,
+ "browser_download_url": f"https://example.invalid/{asset_name}",
+ },
+ ],
+ }
+
+
+@pytest.mark.skipif(
+ direct_linux_release_plan is None,
+ reason = "simple-policy dispatcher not present on this branch",
+)
+def test_simple_policy_plans_lemonade_for_rocm_host():
+ host = _make_rocm_host("gfx1151")
+ with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
+ plan = direct_linux_release_plan(
+ _stub_unsloth_release(),
+ host,
+ "unslothai/llama.cpp",
+ "latest",
+ )
+ assert plan is not None, "ROCm host should not be skipped by simple-policy planner"
+ kinds = [a.install_kind for a in plan.attempts]
+ assert (
+ "linux-rocm" in kinds
+ ), f"simple-policy planner did not include a lemonade ROCm attempt; got {kinds}"
+ rocm_attempt = next(a for a in plan.attempts if a.install_kind == "linux-rocm")
+ assert rocm_attempt.source_label == "lemonade"
+ assert "gfx1151" in rocm_attempt.name
+
+
+@pytest.mark.skipif(
+ direct_upstream_release_plan is None,
+ reason = "simple-policy dispatcher not present on this branch",
+)
+def test_simple_policy_plans_lemonade_for_windows_hip_host():
+ host = _make_rocm_host("gfx1151", windows = True)
+ release = {
+ "tag_name": "b9022",
+ "name": "b9022",
+ "assets": [],
+ }
+ with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
+ plan = direct_upstream_release_plan(
+ release, host, "ggml-org/llama.cpp", "latest"
+ )
+ assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt"
+ kinds = [a.install_kind for a in plan.attempts]
+ assert (
+ "windows-hip" in kinds
+ ), f"simple-policy planner did not include a lemonade HIP attempt; got {kinds}"
+
+
+@pytest.mark.skipif(
+ direct_upstream_release_plan is None,
+ reason = "simple-policy dispatcher not present on this branch",
+)
+def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable():
+ """If lemonade returns None (e.g. gfx999 or transient API failure), the planner
+ must still include the upstream HIP asset rather than silently downgrading to CPU."""
+ host = _make_rocm_host("gfx999", windows = True)
+ hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip"
+ release = {
+ "tag_name": "b9022",
+ "name": "b9022",
+ "assets": [
+ {
+ "name": hip_asset,
+ "browser_download_url": f"https://example.invalid/{hip_asset}",
+ },
+ ],
+ }
+ plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
+ assert plan is not None
+ kinds = [a.install_kind for a in plan.attempts]
+ assert (
+ "windows-hip" in kinds
+ ), f"upstream HIP asset not included as fallback; got {kinds}"
+ hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip")
+ assert hip_attempt.source_label == "upstream"
+
+
+# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ──
+
+
+def test_lemonade_release_api_url_pinned_tag():
+ """A pinned llama_tag must produce the /releases/tags/ URL."""
+ assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262")
+ assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest")
+ assert _mod._lemonade_release_api_for("").endswith("/releases/latest")
+
+
+def test_lemonade_release_api_url_encodes_tag():
+ """Unexpected slashes / hashes in the tag must be URL-encoded so the URL
+ cannot be reshaped (defence in depth -- tags should already be sanitised
+ upstream)."""
+ url = _mod._lemonade_release_api_for("b1260/../latest")
+ assert "/releases/tags/b1260%2F..%2Flatest" in url
+ assert "//latest" not in url.split("/releases/tags/", 1)[1]
+
+
+def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch):
+ """UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver."""
+ monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1")
+ host = _make_rocm_host("gfx1151")
+ res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
+ assert res is None
+
+
+def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
+ """If the GitHub API response somehow contained an off-host download URL,
+ the resolver must refuse to use it (lemonade assets are not in the
+ approved-hash manifest)."""
+ bad_release = {
+ "tag_name": _STUB_TAG,
+ "assets": [
+ {
+ "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
+ "browser_download_url": "https://attacker.invalid/llama.zip",
+ },
+ ],
+ }
+ host = _make_rocm_host("gfx1151")
+ with patch.object(_mod, "fetch_json", return_value = bad_release):
+ res = resolve_lemonade_rocm_choice(
+ host, "ubuntu", "linux-rocm", llama_tag = "latest"
+ )
+ assert res is None
+
+
+def test_lemonade_resolver_rejects_http_scheme():
+ assert not _mod._is_trusted_github_release_url(
+ "http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip",
+ "lemonade-sdk/llamacpp-rocm",
+ )
+
+
+def test_lemonade_resolver_accepts_github_cdn():
+ # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix.
+ assert _mod._is_trusted_github_release_url(
+ "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x",
+ "lemonade-sdk/llamacpp-rocm",
+ )
+
+
+def test_lemonade_resolver_rejects_arbitrary_cdn_path():
+ # A CDN URL without the release-asset path prefix must be rejected.
+ assert not _mod._is_trusted_github_release_url(
+ "https://objects.githubusercontent.com/abc/def",
+ "lemonade-sdk/llamacpp-rocm",
+ )
+
+
+def test_lemonade_resolver_accepts_release_path():
+ url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip"
+ assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm")
+
+
+def test_lemonade_resolver_rejects_wrong_repo():
+ """A github.com release URL for a different repo must be rejected."""
+ assert not _mod._is_trusted_github_release_url(
+ "https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip",
+ "lemonade-sdk/llamacpp-rocm",
+ )
+
+
+def test_lemonade_resolver_rejects_empty_browser_download_url():
+ """An asset entry with an empty browser_download_url must fall through."""
+ release = {
+ "tag_name": _STUB_TAG,
+ "assets": [
+ {
+ "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
+ "browser_download_url": "",
+ },
+ ],
+ }
+ host = _make_rocm_host("gfx1151")
+ with patch.object(_mod, "fetch_json", return_value = release):
+ res = resolve_lemonade_rocm_choice(
+ host, "ubuntu", "linux-rocm", llama_tag = "latest"
+ )
+ assert res is None
+
+
+def test_lemonade_runtime_patterns_include_hip_runtime():
+ """linux-rocm overlay must use a broad lib glob to catch all bundled .so files.
+
+ Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp,
+ ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob
+ avoids having to enumerate every transitive dependency by name.
+ """
+ from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice
+
+ choice = AssetChoice(
+ repo = "lemonade-sdk/llamacpp-rocm",
+ tag = "b1262",
+ name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip",
+ url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip",
+ source_label = "lemonade",
+ install_kind = "linux-rocm",
+ )
+ pats = runtime_patterns_for_choice(choice)
+ # The broad glob must be present so every .so in the lemonade bundle
+ # (including transitive deps added in future ROCm releases) gets overlaid.
+ assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}"
+
+
+_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None)
+
+
+@pytest.mark.skipif(
+ _pick_rocm_gfx_target is None,
+ reason = "_pick_rocm_gfx_target not present on this branch",
+)
+def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch):
+ """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES;
+ on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
+ # Two GPUs; rocminfo reports each token twice (as in the real tool output).
+ probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
+ monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
+ monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
+ assert _pick_rocm_gfx_target(probe_out) == "gfx1100"
+
+
+@pytest.mark.skipif(
+ _pick_rocm_gfx_target is None,
+ reason = "_pick_rocm_gfx_target not present on this branch",
+)
+def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch):
+ """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None."""
+ probe_out = "gfx1151\ngfx1100"
+ monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
+ monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1")
+ assert _pick_rocm_gfx_target(probe_out) is None
+
+
+@pytest.mark.skipif(
+ _pick_rocm_gfx_target is None,
+ reason = "_pick_rocm_gfx_target not present on this branch",
+)
+def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
+ """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must
+ return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the
+ two gfx1100 entries into one and making index 2 out of range."""
+ # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
+ # Each GPU gets its own Agent section with a few token mentions.
+ probe_out = (
+ "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n"
+ "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n"
+ "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n"
+ )
+ monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
+ monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2")
+ assert _pick_rocm_gfx_target(probe_out) == "gfx1151"
diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py
index 1ea76edd15..6fe5372147 100644
--- a/studio/backend/tests/test_llama_cpp_context_fit.py
+++ b/studio/backend/tests/test_llama_cpp_context_fit.py
@@ -84,6 +84,7 @@ _httpx_stub.Client = type(
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend
+from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
# ---------------------------------------------------------------------------
@@ -131,6 +132,7 @@ def _drive(
native_ctx = 131072,
kv_per_token_bytes = 325_000,
can_estimate_kv = True,
+ extra_args = None,
):
"""Drive the post-metadata portion of load_model with stubbed inputs.
@@ -148,11 +150,16 @@ def _drive(
inst._can_estimate_kv = lambda: can_estimate_kv
context_length = inst._context_length
+ # Use the production helper instead of reimplementing the conditional
+ # locally; reimplementing makes the test pass for the test's own logic
+ # rather than production's, and silent drift won't be caught.
+ ctx_override = parse_ctx_override(extra_args)
+ requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
- effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0)
+ effective_ctx = requested_ctx if requested_ctx > 0 else (context_length or 0)
max_available_ctx = context_length or effective_ctx
- if n_ctx > 0:
- effective_ctx = n_ctx
+ if requested_ctx > 0:
+ effective_ctx = requested_ctx
elif context_length is not None:
effective_ctx = context_length
else:
@@ -161,7 +168,7 @@ def _drive(
max_available_ctx = context_length or effective_ctx
gpu_indices, use_fit = None, True
- explicit_ctx = n_ctx > 0
+ explicit_ctx = requested_ctx > 0
if gpus and inst._can_estimate_kv() and effective_ctx > 0:
native_ctx_for_cap = context_length or effective_ctx
@@ -236,6 +243,7 @@ def _drive(
"gpu_indices": gpu_indices,
"max_available_ctx": max_available_ctx,
"original_ctx": original_ctx,
+ "ctx_override": ctx_override,
}
@@ -349,6 +357,48 @@ class TestExplicitCtxRespectsUser:
assert plan["c_arg"] == 2048
+# ---------------------------------------------------------------------------
+# Pass-through --ctx-size participates in context fit (#5676).
+# ---------------------------------------------------------------------------
+
+
+class TestExtraArgsCtxOverride:
+ def test_ctx_size_extra_honored_over_auto(self):
+ plan = _drive(
+ n_ctx = 0,
+ model_gib = 131,
+ gpus = [(0, 97_000)],
+ native_ctx = 196608,
+ extra_args = ["--ctx-size", "128000"],
+ )
+ assert plan["ctx_override"] == 128000
+ assert plan["original_ctx"] == 128000
+ assert plan["c_arg"] == 128000
+ assert plan["use_fit"] is True
+
+ def test_ctx_size_short_alias_honored_over_auto(self):
+ plan = _drive(
+ n_ctx = 0,
+ model_gib = 131,
+ gpus = [(0, 97_000)],
+ native_ctx = 196608,
+ extra_args = ["-c", "128000"],
+ )
+ assert plan["c_arg"] == 128000
+ assert plan["use_fit"] is True
+
+ def test_ctx_size_extra_wins_over_first_class_field(self):
+ plan = _drive(
+ n_ctx = 4096,
+ model_gib = 8,
+ gpus = [(0, 24_000)],
+ native_ctx = 131072,
+ extra_args = ["--ctx-size", "128000"],
+ )
+ assert plan["original_ctx"] == 128000
+ assert plan["c_arg"] == 128000
+
+
# ---------------------------------------------------------------------------
# Non-regression: fittable + auto still auto-picks largest fitting ctx
# ---------------------------------------------------------------------------
diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
new file mode 100644
index 0000000000..e647ff2c7d
--- /dev/null
+++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Tests for LlamaCppBackend._classify_llama_start_failure.
+
+When llama-server exits before becoming healthy, load_model turns its
+captured stdout/stderr into a user-facing reason. A diffusion / image
+GGUF (FLUX, Qwen-Image, ...) is a valid file with plenty of memory, so
+the generic "invalid file or out of memory" message is actively
+misleading (issue #5842). These tests pin the classification.
+"""
+
+from __future__ import annotations
+
+import sys
+import types as _types
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+ sys.path.insert(0, _BACKEND_DIR)
+
+# Match the stubbing pattern in sibling tests so the module imports in a
+# lightweight env without fastapi.
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+# Give the structlog stub a real get_logger: a bare ModuleType poisons
+# sys.modules for later tests that call structlog.get_logger at import time.
+_structlog_stub = _types.ModuleType("structlog")
+_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger(
+ "structlog"
+)
+sys.modules.setdefault("structlog", _structlog_stub)
+if not hasattr(sys.modules["structlog"], "get_logger"):
+ sys.modules["structlog"].get_logger = _structlog_stub.get_logger
+
+from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
+
+_classify = LlamaCppBackend._classify_llama_start_failure
+
+# Real llama-server failure lines (lower-cased downstream anyway).
+_QWEN_IMAGE_OUT = (
+ "load_model: loading model 'qwen-image-edit-2511-Q4_K_M.gguf'\n"
+ "llama_model_load: error loading model: unknown model architecture: 'qwen_image'\n"
+ "llama_model_load_from_file_impl: failed to load model"
+)
+_OOM_OUT = (
+ "ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on "
+ "device 0: cudaMalloc failed: out of memory"
+)
+
+
+class TestDiffusionArchitectures:
+ def test_qwen_image_routes_to_images_page(self):
+ msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image")
+ assert "diffusion" in msg.lower()
+ assert "Images page" in msg
+ assert "qwen_image" in msg
+ # Must NOT keep blaming memory / file validity.
+ assert "out of memory" not in msg.lower()
+ assert "enough memory" not in msg.lower()
+
+ # Parametrize over the production set so new arches are auto-covered.
+ @pytest.mark.parametrize("arch", sorted(LlamaCppBackend._DIFFUSION_ARCHES))
+ def test_every_diffusion_arch_is_recognised(self, arch):
+ out = f"error loading model: unknown model architecture: '{arch}'"
+ msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}")
+ assert "diffusion" in msg.lower()
+ assert "Images page" in msg
+ assert arch in msg
+
+
+class TestUnsupportedNonDiffusionArchitecture:
+ def test_unknown_llm_arch_says_unsupported_not_oom(self):
+ out = "error loading model: unknown model architecture: 'some_new_llm'"
+ msg = _classify(out, "/models/x.gguf", "local/x")
+ assert "some_new_llm" in msg
+ assert "architecture" in msg.lower()
+ # Specific, not the misleading memory message.
+ assert "enough memory" not in msg.lower()
+ assert "diffusion" not in msg.lower()
+
+ # Exact match: a chat arch merely containing a diffusion token (wan,
+ # sd1, flux, ...) must not be routed to the Images page.
+ @pytest.mark.parametrize(
+ "arch",
+ [
+ "taiwan", # contains "wan"
+ "swan_llm", # contains "wan"
+ "fluxion", # contains "flux"
+ "sd1234", # contains "sd1"
+ "sd3_chat", # contains "sd3"
+ "aura2_text", # contains "aura"
+ "cosmos_reason", # contains "cosmos"
+ "qwen_image_text", # contains "qwen_image"
+ ],
+ )
+ def test_arch_containing_diffusion_token_is_not_misrouted(self, arch):
+ out = f"error loading model: unknown model architecture: '{arch}'"
+ msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}")
+ assert arch in msg
+ assert "does not support" in msg.lower()
+ assert "diffusion" not in msg.lower()
+ assert "Images page" not in msg
+
+
+class TestOllamaAndFallback:
+ _OLLAMA_GGUF = (
+ f"/home/u/.ollama{__import__('os').sep}ollama_links"
+ f"{__import__('os').sep}m.gguf"
+ )
+
+ def test_ollama_compat_message_still_works(self):
+ out = "llama_model_load: error loading model: key not found"
+ msg = _classify(out, self._OLLAMA_GGUF, "ollama/llama3")
+ assert "Ollama" in msg
+
+ def test_ollama_unknown_arch_keeps_ollama_guidance(self):
+ # Ollama + non-diffusion unknown arch keeps the Ollama hint, not the
+ # generic llama.cpp "unsupported" message.
+ out = "error loading model: unknown model architecture: 'some_new_llm'"
+ msg = _classify(out, self._OLLAMA_GGUF, "ollama/some-new")
+ assert "Ollama" in msg
+ assert "directly through Ollama" in msg
+ assert "does not support" not in msg.lower()
+
+ def test_ollama_diffusion_arch_still_routes_to_images(self):
+ # Diffusion routing wins over the Ollama hint.
+ out = "error loading model: unknown model architecture: 'flux'"
+ msg = _classify(out, self._OLLAMA_GGUF, "ollama/flux")
+ assert "diffusion" in msg.lower()
+ assert "Images page" in msg
+
+ def test_generic_oom_keeps_memory_message(self):
+ msg = _classify(_OOM_OUT, "/models/big.gguf", "local/big")
+ assert "enough memory" in msg.lower()
+ assert "diffusion" not in msg.lower()
+
+ def test_empty_output_is_safe(self):
+ msg = _classify("", None, None)
+ assert "llama-server failed to start" in msg
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index 68a1c870fb..2f7431497d 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -3,21 +3,38 @@
"""Unit tests for the llama-server pass-through args validator.
-The validator is the security boundary between user-supplied CLI / HTTP
-input and the llama-server subprocess command. These tests pin the
-denylist behavior so the boundary doesn't quietly regress when new
-managed flags are added.
+The validator is the boundary between user CLI/HTTP input and the
+llama-server subprocess. These tests pin denylist behaviour so it
+doesn't quietly regress when new managed flags are added.
"""
from __future__ import annotations
+import importlib.util
+import re
+from pathlib import Path
+
import pytest
-from core.inference.llama_server_args import (
- is_managed_flag,
- strip_shadowing_flags,
- validate_extra_args,
+# Load llama_server_args.py directly so this test doesn't drag in the
+# full backend chain (fastapi / structlog / loggers / utils.hardware)
+# via core/inference/__init__.py. The validator is intentionally
+# dependency-free and unit-tests should reflect that.
+_LSA_PATH = (
+ Path(__file__).resolve().parent.parent
+ / "core"
+ / "inference"
+ / "llama_server_args.py"
)
+_spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH)
+_lsa = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_lsa)
+is_managed_flag = _lsa.is_managed_flag
+parse_cache_override = _lsa.parse_cache_override
+parse_ctx_override = _lsa.parse_ctx_override
+resolve_cache_type_kv = _lsa.resolve_cache_type_kv
+strip_shadowing_flags = _lsa.strip_shadowing_flags
+validate_extra_args = _lsa.validate_extra_args
# ── Pass-through (allowed) ───────────────────────────────────────────
@@ -60,13 +77,12 @@ from core.inference.llama_server_args import (
# Reasoning controls
["--reasoning-format", "deepseek"],
["-rea", "auto"],
- # Soft-managed flags the user may want to override on the CLI;
- # llama.cpp's last-wins parsing means these win over Studio's
- # auto-set version.
+ # Soft-managed: user-supplied flags last-wins-override Studio's
+ # auto-set version. --parallel / -np / --n-parallel are NOT
+ # here -- they're hard-denied (KV-cache + slot count would
+ # desync). Use `unsloth studio run --parallel N` instead.
["-c", "131072"],
["--ctx-size", "8192"],
- ["--parallel", "1"],
- ["-np", "8"],
["--flash-attn", "off"],
["-fa", "on"],
["--no-context-shift"],
@@ -99,8 +115,7 @@ def test_value_with_equals_form_passes_through():
def test_non_flag_token_passes_through():
- # A bare positional value (not preceded by a flag) is preserved
- # verbatim. llama-server may reject it, but that's not our job.
+ # Bare positionals are passed through; llama-server can reject them.
assert validate_extra_args(["foo"]) == ["foo"]
@@ -110,18 +125,33 @@ def test_non_flag_token_passes_through():
@pytest.mark.parametrize(
"denied",
[
- # Model identity
+ # Parallel slots -- owned by the typer --parallel flag.
+ "-np",
+ "--parallel",
+ "--n-parallel",
+ # Model identity (every alias; bumping llama.cpp must keep
+ # every form rejected, not just the long).
"-m",
"--model",
+ "-mu",
+ "--model-url",
+ "-dr",
+ "--docker-repo",
"-hf",
"-hfr",
"--hf-repo",
"-hff",
"--hf-file",
+ "-hfv",
+ "-hfrv",
+ "--hf-repo-v",
+ "-hffv",
+ "--hf-file-v",
"-hft",
"--hf-token",
"-mm",
"--mmproj",
+ "-mmu",
"--mmproj-url",
# Networking (Studio binds + proxies)
"--host",
@@ -134,11 +164,28 @@ def test_non_flag_token_passes_through():
"--api-key-file",
"--ssl-key-file",
"--ssl-cert-file",
- # Single-model server
+ # Single-model server (legacy --webui + current --ui group)
"--webui",
"--no-webui",
+ "--ui",
+ "--no-ui",
+ "--ui-config",
+ "--ui-config-file",
+ "--ui-mcp-proxy",
+ "--no-ui-mcp-proxy",
"--models-dir",
+ "--models-preset",
"--models-max",
+ "--models-autoload",
+ "--no-models-autoload",
+ # Server-mode flips: --embedding / --rerank would restrict
+ # llama-server to those endpoints and break Studio's chat hop.
+ "--embedding",
+ "--embeddings",
+ "--rerank",
+ "--reranking",
+ # llama-server's own --tools clashes with Studio's tool policy.
+ "--tools",
],
)
def test_denylist_rejects_all_aliases(denied):
@@ -146,14 +193,65 @@ def test_denylist_rejects_all_aliases(denied):
validate_extra_args([denied, "value"])
+@pytest.mark.parametrize(
+ "args,offending",
+ [
+ # Pass-through --parallel would last-wins-override the real
+ # slot count while Studio's KV-cache fit + llama_parallel_slots
+ # stay at the typer value -- plan vs. process disagree.
+ (["--parallel", "8"], "--parallel"),
+ (["--parallel=8"], "--parallel"),
+ (["--n-parallel", "16"], "--n-parallel"),
+ (["--n-parallel=16"], "--n-parallel"),
+ (["-np", "32"], "-np"),
+ # Attached short form: Click clusters it CLI-side; HTTP /load
+ # with `["-np8"]` must still resolve to managed.
+ (["-np8"], "-np"),
+ (["-np64"], "-np"),
+ # Out-of-range values that would bypass the typer 1..64 guard.
+ (["--parallel", "999"], "--parallel"),
+ (["-np", "0"], "-np"),
+ (["-np999"], "-np"),
+ # Signed attached forms; `-np-1` must not slip past.
+ (["-np-1"], "-np"),
+ (["-np+1"], "-np"),
+ ],
+)
+def test_parallel_flags_are_managed(args, offending):
+ with pytest.raises(ValueError, match = re.escape(offending)):
+ validate_extra_args(args)
+
+
def test_denylist_rejects_equals_form():
with pytest.raises(ValueError, match = "--port"):
validate_extra_args(["--port=9000"])
+@pytest.mark.parametrize(
+ "padded",
+ [" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"],
+)
+def test_denylist_rejects_whitespace_padded_forms(padded):
+ # `_flag_name` trims whitespace before lookup; otherwise a trailing
+ # space could slip a managed flag past the boundary.
+ with pytest.raises(ValueError, match = "parallel|np"):
+ validate_extra_args([padded, "8"])
+
+
+@pytest.mark.parametrize(
+ "attached",
+ ["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"],
+)
+def test_denylist_rejects_np_with_digit_prefix_and_junk(attached):
+ # Backend `_flag_name` must classify the same forms the CLI
+ # rewriter expands, else HTTP /load could smuggle `-np8x` through.
+ with pytest.raises(ValueError, match = "np"):
+ validate_extra_args([attached])
+
+
def test_denylist_rejects_short_form_when_long_is_denied():
- # -m is the short form of the hard-denied --model; rejecting only
- # the long form would leave a trivial bypass.
+ # `-m` is the short form of --model; rejecting only the long
+ # form would leave a trivial bypass.
with pytest.raises(ValueError, match = "-m"):
validate_extra_args(["-m", "/some/other/path.gguf"])
@@ -165,9 +263,7 @@ def test_denylist_message_names_offending_flag():
def test_first_denied_flag_short_circuits():
- # Validation stops at the first denied flag; later denied flags
- # in the same call don't matter for behaviour, but the message
- # should name the first one we hit.
+ # Validation stops at the first denied flag; the message names it.
with pytest.raises(ValueError, match = "--port"):
validate_extra_args(["--port", "1", "--host", "x"])
@@ -177,8 +273,7 @@ def test_first_denied_flag_short_circuits():
@pytest.mark.parametrize("value", ["-1", "-0.5", "-42", "-.5"])
def test_negative_number_value_is_not_flag(value):
- # ``--seed -1`` is a value, not a flag. Validator must not try
- # to look up "-1" in the denylist.
+ # `--seed -1`: the -1 is a value, not a flag.
assert validate_extra_args(["--seed", value]) == ["--seed", value]
@@ -190,6 +285,15 @@ def test_is_managed_flag_true_for_denied():
assert is_managed_flag("--api-key") is True
assert is_managed_flag("-m") is True
assert is_managed_flag("--model") is True
+ # Parallel slots owned by the typer --parallel flag.
+ assert is_managed_flag("--parallel") is True
+ assert is_managed_flag("--n-parallel") is True
+ assert is_managed_flag("-np") is True
+ # Normalised forms must classify like the canonical token so
+ # is_managed_flag filtering stays in sync with validate_extra_args.
+ assert is_managed_flag("-np8") is True
+ assert is_managed_flag("--parallel=8") is True
+ assert is_managed_flag("--port=9000") is True
def test_is_managed_flag_false_for_pass_through():
@@ -199,7 +303,6 @@ def test_is_managed_flag_false_for_pass_through():
# Soft-managed flags pass through (last-wins override)
assert is_managed_flag("-c") is False
assert is_managed_flag("--ctx-size") is False
- assert is_managed_flag("--parallel") is False
assert is_managed_flag("--flash-attn") is False
assert is_managed_flag("-ngl") is False
assert is_managed_flag("--threads") is False
@@ -231,8 +334,8 @@ def test_strip_shadowing_flags_keeps_context_when_not_requested():
def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled():
- # Caller did not supply chat_template_override; the inherited
- # --chat-template-file must survive the strip.
+ # No chat_template_override supplied; inherited
+ # --chat-template-file must survive.
out = strip_shadowing_flags(
["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"],
strip_context = True,
@@ -282,7 +385,7 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
def test_strip_shadowing_flags_drops_mtp_flags_when_requested():
- # MTP / draft-mtp flags must be stripped when speculative_type is re-applied.
+ # MTP / draft-mtp flags must drop when speculative_type re-applies.
out = strip_shadowing_flags(
[
"--spec-type",
@@ -310,9 +413,88 @@ def test_is_managed_flag_false_for_mtp_pass_through():
assert is_managed_flag("--spec-ngram-mod-n-max") is False
+# ── parse_ctx_override ───────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+ "args,expected",
+ [
+ (None, None),
+ ([], None),
+ (["--top-k", "20"], None),
+ (["--ctx-size", "128000"], 128000),
+ (["--ctx-size=128000"], 128000),
+ (["-c", "128000"], 128000),
+ (["-c=128000"], 128000),
+ (["-c", "4096", "--ctx-size", "128000"], 128000),
+ ],
+)
+def test_parse_ctx_override(args, expected):
+ assert parse_ctx_override(args) == expected
+
+
+@pytest.mark.parametrize(
+ "args",
+ [
+ ["--ctx-size"],
+ ["--ctx-size", "--top-k"],
+ ["--ctx-size", "abc"],
+ ["--ctx-size=abc"],
+ ["-c", "-1"],
+ ],
+)
+def test_parse_ctx_override_rejects_malformed_values(args):
+ with pytest.raises(ValueError, match = "ctx-size|'-c'"):
+ parse_ctx_override(args)
+
+
+def test_validate_extra_args_rejects_malformed_ctx_override():
+ with pytest.raises(ValueError, match = "ctx-size"):
+ validate_extra_args(["--ctx-size", "abc"])
+
+
+# ── parse_cache_override ─────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+ "args,expected",
+ [
+ (None, None),
+ ([], None),
+ (["--top-k", "20"], None),
+ (["--cache-type-k", "q8_0"], "q8_0"),
+ (["-ctk", "q4_0"], "q4_0"),
+ (["-ctv", "q4_0"], "q4_0"),
+ (["--cache-type-k=q4_0"], "q4_0"),
+ (["-ctk", "f16", "-ctk", "q8_0"], "q8_0"),
+ ],
+)
+def test_parse_cache_override(args, expected):
+ assert parse_cache_override(args) == expected
+
+
+@pytest.mark.parametrize(
+ "args",
+ [
+ ["-ctk"],
+ ["-ctk", "-c", "4096"],
+ ],
+)
+def test_parse_cache_override_rejects_malformed_values(args):
+ with pytest.raises(ValueError, match = "cache-type|'-ctk'"):
+ parse_cache_override(args)
+
+
+def test_resolve_cache_type_kv_uses_override_when_present():
+ assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0"
+
+
+def test_resolve_cache_type_kv_uses_fallback_without_override():
+ assert resolve_cache_type_kv(["--top-k", "20"], "f16") == "f16"
+
+
def test_strip_shadowing_flags_boolean_does_not_consume_next_token():
- # --spec-default is a boolean shadowing flag; the value-skipping
- # heuristic must skip just the flag, not the following positional.
+ # `--spec-default` is boolean; drop just the flag, keep the next token.
out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True)
assert out == ["ngram-mod"]
@@ -343,8 +525,8 @@ def test_strip_shadowing_flags_handles_empty_input():
def test_strip_shadowing_flags_defaults_strip_everything():
- # The route's already-loaded comparator calls strip_shadowing_flags
- # with no kwargs to detect ANY shadowing flag in stored extras.
+ # The route's already-loaded comparator calls with no kwargs to
+ # detect ANY shadowing flag in stored extras.
out = strip_shadowing_flags(
["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"]
)
diff --git a/studio/backend/tests/test_log_filter_no_truncation.py b/studio/backend/tests/test_log_filter_no_truncation.py
index d78643f5b9..d9a6e2bc4a 100644
--- a/studio/backend/tests/test_log_filter_no_truncation.py
+++ b/studio/backend/tests/test_log_filter_no_truncation.py
@@ -2,27 +2,11 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
-Regression tests for studio.backend.loggers.handlers.filter_sensitive_data.
+Regression tests for loggers.handlers.filter_sensitive_data.
-Context: filter_sensitive_data was originally written with a base64-detection
-heuristic that truncated any string >100 chars containing ',' or '/' down to
-20 chars + '...'. The block was dormant until PR #5246 wired the processor
-into the structlog chain to redact native-path leases. Once active, the
-heuristic ate normal log lines emitted by llama_cpp_backend (GGUF size
-summary, mmproj selection, the full llama-server command line) and any
-exception traceback that happened to contain a file path.
-
-These tests pin two properties:
-
-1. Long, comma- or slash-bearing log messages flow through filter_sensitive_data
- unchanged. The exact strings exercised match the call sites at
- studio/backend/core/inference/llama_cpp.py:2117, :2283, and :2312 that
- were truncated in the original bug report.
-
-2. PR #5246's native-path lease redaction still fires for both the inline
- ``native_path_lease=...`` regex form and the ``nativePathLease`` dict-key
- form. This guards against future regressions that strip redaction along
- with the truncation block.
+Pins two properties: (1) long strings with commas/slashes pass through
+unchanged (the base64-truncation heuristic from PR #5246 was too aggressive),
+and (2) native-path lease redaction still fires for both inline and dict-key forms.
"""
from loggers.handlers import filter_sensitive_data
diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py
new file mode 100644
index 0000000000..10a6eb012b
--- /dev/null
+++ b/studio/backend/tests/test_mcp_servers.py
@@ -0,0 +1,632 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import pytest
+from fastapi import HTTPException
+
+from storage import mcp_servers_db
+
+
+def _reset_db(tmp_path, monkeypatch):
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
+
+
+# ── storage: mcp_servers_db ─────────────────────────────────────────
+
+
+def test_create_and_get_server(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(
+ id = "srv1",
+ display_name = "GitHub",
+ url = "https://example.com/mcp",
+ headers_json = '{"Authorization": "Bearer x"}',
+ is_enabled = True,
+ use_oauth = False,
+ )
+ row = mcp_servers_db.get_server("srv1")
+ assert row["id"] == "srv1"
+ assert row["display_name"] == "GitHub"
+ assert row["url"] == "https://example.com/mcp"
+ assert row["headers_json"] == '{"Authorization": "Bearer x"}'
+ assert row["is_enabled"] == 1
+ assert row["use_oauth"] == 0
+
+
+def test_list_servers_ordered_by_created_at(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(id = "a", display_name = "A", url = "https://a/m")
+ mcp_servers_db.create_server(id = "b", display_name = "B", url = "https://b/m")
+ rows = mcp_servers_db.list_servers()
+ assert [r["id"] for r in rows] == ["a", "b"]
+
+
+def test_update_server_coerces_bools(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
+ assert mcp_servers_db.update_server(
+ "srv1", {"is_enabled": False, "use_oauth": True}
+ )
+ row = mcp_servers_db.get_server("srv1")
+ assert row["is_enabled"] == 0
+ assert row["use_oauth"] == 1
+
+
+def test_update_server_empty_changes_returns_false(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
+ assert mcp_servers_db.update_server("srv1", {}) is False
+
+
+def test_delete_server_roundtrip(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
+ assert mcp_servers_db.delete_server("srv1") is True
+ assert mcp_servers_db.delete_server("srv1") is False
+ assert mcp_servers_db.get_server("srv1") is None
+
+
+# ── routes/mcp_servers: pure helpers ────────────────────────────────
+
+
+def test_validate_url_accepts_http_and_https():
+ from routes.mcp_servers import _validate_url
+
+ assert _validate_url("http://example.com/mcp") == "http://example.com/mcp"
+ assert _validate_url("https://example.com/mcp") == "https://example.com/mcp"
+ assert _validate_url(" https://example.com/mcp ") == "https://example.com/mcp"
+
+
+@pytest.mark.parametrize("bad", ["", " ", "ftp://x", "http://", "noscheme.com"])
+def test_validate_url_rejects_bad(bad):
+ from routes.mcp_servers import _validate_url
+
+ with pytest.raises(HTTPException) as exc:
+ _validate_url(bad)
+ assert exc.value.status_code == 400
+
+
+def test_normalize_headers():
+ from routes.mcp_servers import _normalize_headers
+
+ assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == {
+ "Auth": "Bearer x"
+ }
+ assert _normalize_headers({"X": 42}) == {"X": "42"}
+ assert _normalize_headers({}) is None
+ assert _normalize_headers(None) is None
+ assert _normalize_headers({" ": "x"}) is None
+
+
+def test_changes_from_payload_tristate_headers():
+ from routes.mcp_servers import _changes_from_payload
+ from models.mcp_servers import McpServerUpdate
+
+ # omitted → key absent
+ assert "headers_json" not in _changes_from_payload(
+ McpServerUpdate(display_name = "x")
+ )
+ # null → stored as None (clear all headers)
+ assert _changes_from_payload(McpServerUpdate(headers = None))["headers_json"] is None
+ # dict → serialised JSON
+ assert (
+ _changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"]
+ == '{"a": "1"}'
+ )
+
+
+# ── core/inference/tools: MCP wiring ────────────────────────────────
+
+
+def test_mcp_specs_skip_oversized_names():
+ from core.inference.tools import _mcp_specs_for_server
+
+ server = {"id": "s" * 30, "display_name": "S"}
+ tools = [
+ {"name": "ok", "description": "fine"},
+ {"name": "x" * 40, "description": "too long"},
+ ]
+ specs = _mcp_specs_for_server(server, tools)
+ assert len(specs) == 1
+ assert specs[0]["function"]["name"].endswith("__ok")
+ assert len(specs[0]["function"]["name"]) <= 64
+
+
+def test_execute_tool_malformed_mcp_name():
+ from core.inference.tools import execute_tool
+
+ out = execute_tool("mcp__no_double_underscore", {})
+ assert out.startswith("Error: malformed MCP tool name")
+
+
+def test_execute_tool_unknown_server(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ from core.inference.tools import execute_tool
+
+ assert (
+ execute_tool("mcp__missing__do_thing", {})
+ == "Error: MCP server 'missing' not found"
+ )
+
+
+def test_execute_tool_disabled_server(tmp_path, monkeypatch):
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(
+ id = "srv1",
+ display_name = "A",
+ url = "https://a/m",
+ is_enabled = False,
+ )
+ from core.inference.tools import execute_tool
+
+ assert (
+ execute_tool("mcp__srv1__do_thing", {})
+ == "Error: MCP server 'srv1' is disabled"
+ )
+
+
+def test_mcp_specs_skip_invalid_openai_function_names():
+ """OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; tools whose
+ names contain '.', '/', spaces, etc. would 400 the whole request."""
+ from core.inference.tools import _mcp_specs_for_server
+
+ server = {"id": "srv", "display_name": "S"}
+ tools = [
+ {"name": "ok"},
+ {"name": "with.dot"},
+ {"name": "weird/slash"},
+ {"name": "has space"},
+ {"name": "good-dash_ok"},
+ ]
+ specs = _mcp_specs_for_server(server, tools)
+ names = {s["function"]["name"] for s in specs}
+ assert {"mcp__srv__ok", "mcp__srv__good-dash_ok"} == names
+
+
+def test_mcp_specs_skip_empty_tool_name():
+ from core.inference.tools import _mcp_specs_for_server
+
+ server = {"id": "srv", "display_name": "S"}
+ specs = _mcp_specs_for_server(server, [{"name": "", "description": "x"}])
+ assert specs == []
+
+
+def test_mcp_specs_drops_duplicate_names():
+ """Same tool name twice from one MCP server -> OpenAI rejects the
+ request as 'duplicates'. Drop the duplicate before forwarding."""
+ from core.inference.tools import _mcp_specs_for_server
+
+ server = {"id": "srv", "display_name": "S"}
+ tools = [{"name": "echo"}, {"name": "echo"}]
+ specs = _mcp_specs_for_server(server, tools)
+ assert len(specs) == 1
+
+
+def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
+ """cancel_event already set before the call -> immediate Error: cancelled
+ without making a network round-trip."""
+ import threading
+ from core.inference import mcp_client
+
+ # Stub _client so the test doesn't need a real MCP server.
+ class _StubClient:
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, *args):
+ return False
+
+ async def call_tool(self, name, args):
+ import asyncio as _asyncio
+
+ await _asyncio.sleep(30) # never finishes within the test
+
+ monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
+
+ cancel = threading.Event()
+ cancel.set()
+ out = mcp_client.call_tool_sync(
+ url = "https://example/mcp",
+ headers = None,
+ name = "slow",
+ args = {},
+ timeout = 30.0,
+ cancel_event = cancel,
+ )
+ assert "cancelled" in out.lower()
+
+
+def test_clear_oauth_tokens_async_no_op_safe(tmp_path, monkeypatch):
+ """clear_oauth_tokens_async on a URL with no stored token must not raise --
+ the delete + update handlers call it best-effort regardless of prior state."""
+ import asyncio
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ from core.inference import mcp_client
+
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+ asyncio.run(mcp_client.clear_oauth_tokens_async("https://example.com/mcp"))
+
+
+def test_delete_server_calls_oauth_cleanup_when_oauth_was_on(tmp_path, monkeypatch):
+ """delete_mcp_server route helper should invoke clear_oauth_tokens_async
+ when the deleted row had use_oauth=true."""
+ import asyncio
+
+ _reset_db(tmp_path, monkeypatch)
+ from core.inference import mcp_client
+
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+ mcp_servers_db.create_server(
+ id = "oauth1",
+ display_name = "GH",
+ url = "https://gh-mcp.example/mcp",
+ is_enabled = True,
+ use_oauth = True,
+ )
+
+ calls: list[str] = []
+
+ async def fake_clear(url):
+ calls.append(url)
+
+ monkeypatch.setattr(mcp_client, "clear_oauth_tokens_async", fake_clear)
+ # Re-import the route's binding through the module so the patch is seen.
+ import routes.mcp_servers as routes_mcp
+
+ monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
+ asyncio.run(routes_mcp.delete_mcp_server("oauth1", current_subject = "u"))
+ assert calls == ["https://gh-mcp.example/mcp"]
+ assert mcp_servers_db.get_server("oauth1") is None
+
+
+def test_delete_server_skips_oauth_cleanup_when_oauth_off(tmp_path, monkeypatch):
+ """No OAuth token cleanup when the deleted server never had OAuth."""
+ import asyncio
+
+ _reset_db(tmp_path, monkeypatch)
+ from core.inference import mcp_client
+ import routes.mcp_servers as routes_mcp
+
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+ mcp_servers_db.create_server(
+ id = "noauth",
+ display_name = "Plain",
+ url = "https://plain/mcp",
+ is_enabled = True,
+ use_oauth = False,
+ )
+ calls: list[str] = []
+
+ async def fake_clear(url):
+ calls.append(url)
+
+ monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
+ asyncio.run(routes_mcp.delete_mcp_server("noauth", current_subject = "u"))
+ assert calls == []
+
+
+def test_update_server_clears_oauth_on_url_change(tmp_path, monkeypatch):
+ """Changing the URL on an OAuth server must drop the old URL's tokens
+ so the new URL doesn't silently inherit credentials."""
+ import asyncio
+
+ _reset_db(tmp_path, monkeypatch)
+ from core.inference import mcp_client
+ from models.mcp_servers import McpServerUpdate
+ import routes.mcp_servers as routes_mcp
+
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+ mcp_servers_db.create_server(
+ id = "s1",
+ display_name = "A",
+ url = "https://old/mcp",
+ is_enabled = True,
+ use_oauth = True,
+ )
+ calls: list[str] = []
+
+ async def fake_clear(url):
+ calls.append(url)
+
+ monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
+ asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1",
+ McpServerUpdate(url = "https://new/mcp"),
+ current_subject = "u",
+ )
+ )
+ assert calls == ["https://old/mcp"]
+ row = mcp_servers_db.get_server("s1")
+ assert row["url"] == "https://new/mcp"
+
+
+def test_update_server_clears_oauth_when_oauth_disabled(tmp_path, monkeypatch):
+ """Flipping use_oauth false must drop the old URL's tokens."""
+ import asyncio
+
+ _reset_db(tmp_path, monkeypatch)
+ from core.inference import mcp_client
+ from models.mcp_servers import McpServerUpdate
+ import routes.mcp_servers as routes_mcp
+
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+ mcp_servers_db.create_server(
+ id = "s1",
+ display_name = "A",
+ url = "https://u/mcp",
+ is_enabled = True,
+ use_oauth = True,
+ )
+ calls: list[str] = []
+
+ async def fake_clear(url):
+ calls.append(url)
+
+ monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
+ asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1",
+ McpServerUpdate(use_oauth = False),
+ current_subject = "u",
+ )
+ )
+ assert calls == ["https://u/mcp"]
+
+
+def test_changes_from_payload_rejects_null_is_enabled():
+ """Explicit null for is_enabled used to hit int(None) -> TypeError 500."""
+ from routes.mcp_servers import _changes_from_payload
+ from models.mcp_servers import McpServerUpdate
+
+ with pytest.raises(HTTPException) as exc:
+ _changes_from_payload(McpServerUpdate(is_enabled = None))
+ assert exc.value.status_code == 400
+
+
+def test_changes_from_payload_rejects_null_use_oauth():
+ """Explicit null for use_oauth used to hit int(None) -> TypeError 500."""
+ from routes.mcp_servers import _changes_from_payload
+ from models.mcp_servers import McpServerUpdate
+
+ with pytest.raises(HTTPException) as exc:
+ _changes_from_payload(McpServerUpdate(use_oauth = None))
+ assert exc.value.status_code == 400
+
+
+def test_test_endpoint_surfaces_url_validation_as_400(tmp_path, monkeypatch):
+ """POST /api/mcp/servers/test must 400 on invalid URL like create/update;
+ previously the same input returned 200 with {"ok": false}."""
+ import asyncio
+
+ _reset_db(tmp_path, monkeypatch)
+ from routes.mcp_servers import test_mcp_server
+ from models.mcp_servers import McpServerTestRequest
+
+ with pytest.raises(HTTPException) as exc:
+ asyncio.run(
+ test_mcp_server(
+ McpServerTestRequest(url = "ftp://nope"),
+ current_subject = "u",
+ )
+ )
+ assert exc.value.status_code == 400
+
+
+def test_tool_xml_parser_handles_hyphenated_parameter_names():
+ """MCP tool schemas commonly use hyphenated property names like
+ `issue-number` / `repo-name`; the XML parser's `` regex
+ dropped those keys. Verify hyphenated parameter names round-trip."""
+ from core.inference.tool_call_parser import parse_tool_calls_from_text
+ import json as _json
+
+ calls = parse_tool_calls_from_text(
+ ""
+ "Bug report"
+ "octocat/hello"
+ ""
+ )
+ assert len(calls) == 1
+ args = _json.loads(calls[0]["function"]["arguments"])
+ assert args == {"issue-title": "Bug report", "repo-name": "octocat/hello"}
+
+
+def test_tool_healing_strip_handles_hyphenated_function_names():
+ """GGUF's core/tool_healing.py has its own copy of the XML strip
+ regex; the round-4 fix to the shared parser missed this file."""
+ from core.tool_healing import strip_tool_call_markup
+
+ out = strip_tool_call_markup(
+ "before "
+ "x after"
+ )
+ assert out == "before after"
+
+
+def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch):
+ """When the model emits a tool call not in the per-request tool list
+ the GGUF agentic loop must refuse to dispatch -- mirroring the
+ safetensors path. Previously execute_tool ran the call regardless."""
+ from core.inference import tools as tools_mod
+
+ captured: list[str] = []
+
+ def fake_execute(name, args, **kw):
+ captured.append(name)
+ return "executed"
+
+ monkeypatch.setattr(tools_mod, "execute_tool", fake_execute)
+
+ # Re-create the allow-list check inline so we can unit-test the
+ # behavior without spinning up llama-server.
+ def _gate(tools_advertised, called_name, args):
+ allowed = {
+ (t.get("function") or {}).get("name")
+ for t in (tools_advertised or [])
+ if (t.get("function") or {}).get("name")
+ }
+ if allowed and called_name not in allowed:
+ return "Error: tool '" + called_name + "' is not enabled"
+ return fake_execute(called_name, args)
+
+ # Built-in not in advertised list -> blocked.
+ out = _gate(
+ [{"function": {"name": "mcp__srv__echo"}}],
+ "terminal",
+ {"command": "echo x"},
+ )
+ assert "not enabled" in out
+ assert captured == []
+ # Tool in advertised list -> runs.
+ out = _gate(
+ [{"function": {"name": "mcp__srv__echo"}}],
+ "mcp__srv__echo",
+ {"text": "hi"},
+ )
+ assert out == "executed"
+ assert captured == ["mcp__srv__echo"]
+
+
+def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch):
+ """cancel_event set BEFORE call_tool_sync runs -> no HTTP request
+ is made. Previously the call task was created before the cancel
+ check, opening a transport that the watcher then had to cancel."""
+ from core.inference import mcp_client
+
+ opened: list[str] = []
+
+ class _StubClient:
+ async def __aenter__(self):
+ opened.append("opened")
+ return self
+
+ async def __aexit__(self, *args):
+ return False
+
+ async def call_tool(self, name, args):
+ return "ran"
+
+ monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
+
+ import threading
+
+ ev = threading.Event()
+ ev.set()
+ out = mcp_client.call_tool_sync(
+ url = "https://example/mcp",
+ headers = None,
+ name = "x",
+ args = {},
+ timeout = 5.0,
+ cancel_event = ev,
+ )
+ assert "cancelled" in out.lower()
+ # The client must NOT have been opened.
+ assert opened == []
+
+
+def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch):
+ """clear_oauth_tokens_async is best-effort; an OAuth constructor
+ failure (e.g. missing fastmcp.client.auth) must not bubble out into
+ a 500 from the delete / update routes."""
+ import asyncio
+ from core.inference import mcp_client
+
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+
+ # Patch the OAuth import path to raise so the entire body fails.
+ class _BoomOAuth:
+ def __init__(self, *a, **kw):
+ raise RuntimeError("simulated")
+
+ import sys as _sys
+
+ fake_mod = type(_sys)("fastmcp.client.auth")
+ fake_mod.OAuth = _BoomOAuth
+ monkeypatch.setitem(_sys.modules, "fastmcp.client.auth", fake_mod)
+ # Must not raise.
+ asyncio.run(mcp_client.clear_oauth_tokens_async("https://x/mcp"))
+
+
+def test_tool_xml_parser_handles_hyphenated_function_names():
+ """MCP tool names are advertised as `mcp__srv__list-issues` (the regex
+ fix allows '-'); the XML tool-call parser must parse them too,
+ otherwise the model can call the tool but Studio cannot dispatch."""
+ from core.inference.tool_call_parser import parse_tool_calls_from_text
+
+ calls = parse_tool_calls_from_text(
+ ""
+ "octocat/hello"
+ ""
+ )
+ assert len(calls) == 1
+ assert calls[0]["function"]["name"] == "mcp__srv__list-issues"
+ import json as _json
+
+ args = _json.loads(calls[0]["function"]["arguments"])
+ assert args == {"repo": "octocat/hello"}
+
+
+def test_tool_xml_strip_handles_hyphenated_function_names():
+ """routes/inference.py:_TOOL_XML_RE must strip a ``
+ block; otherwise hyphenated MCP tool-call XML leaks into chat history."""
+ import re as _re
+ from pathlib import Path
+
+ src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text()
+ m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL)
+ assert m, "could not extract _TOOL_XML_RE"
+ ns: dict = {"_re": _re}
+ exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns)
+ rx = ns["_TOOL_XML_RE"]
+ stripped = rx.sub(
+ "",
+ "before "
+ "x after",
+ )
+ assert stripped == "before after"
+
+
+def test_safetensors_agentic_empty_allowlist_still_means_allow_all():
+ """Document existing contract: at the safetensors_agentic layer,
+ tools=[] is still treated as "no constraint" (so existing callers
+ work unchanged). The real fix for the MCP-only-no-discovery case
+ lives at the route level in inference.py, which refuses to enter
+ use_tools when the resolved tool list is empty."""
+ import threading
+ from core.inference.safetensors_agentic import run_safetensors_tool_loop
+
+ calls: list[str] = []
+
+ def fake_execute(name, args, **kw):
+ calls.append(name)
+ return "ran"
+
+ iteration = {"n": 0}
+
+ def fake_single_turn(messages):
+ iteration["n"] += 1
+ if iteration["n"] == 1:
+ txt = '{"name":"python","arguments":{"code":"1"}} '
+ buf = ""
+ for ch in txt:
+ buf += ch
+ yield buf
+ else:
+ yield "done"
+
+ list(
+ run_safetensors_tool_loop(
+ single_turn = fake_single_turn,
+ messages = [{"role": "user", "content": "x"}],
+ tools = [],
+ execute_tool = fake_execute,
+ cancel_event = threading.Event(),
+ max_tool_iterations = 1,
+ )
+ )
+ # Empty allow-list = run anything (preserved contract).
+ assert calls == [("python", {"code": "1"})] or len(calls) >= 1
diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py
new file mode 100644
index 0000000000..e980e6a057
--- /dev/null
+++ b/studio/backend/tests/test_mcp_stdio_improvements.py
@@ -0,0 +1,236 @@
+"""Tests for the proposed PR #5863 improvements.
+
+Covers: _client() self-gating + keep_alive, OAuth normalised off for stdio
+(create + update), env/header dropped on a transport-type switch, and the
+backend rejecting a command whose first token is a URL scheme.
+
+Run from studio/backend: python -m pytest tests/test_mcp_stdio_improvements.py -q
+"""
+
+import asyncio
+
+import pytest
+from fastapi import HTTPException
+
+from core.inference import mcp_client
+from storage import mcp_servers_db
+
+
+def _reset_db(tmp_path, monkeypatch):
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
+
+
+def _enable(monkeypatch):
+ monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
+
+
+def _disable(monkeypatch):
+ monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
+
+
+# ── P1: _client() self-gates the stdio sink ─────────────────────────
+
+
+def test_client_refuses_stdio_when_disabled(monkeypatch):
+ _disable(monkeypatch)
+ with pytest.raises(PermissionError):
+ mcp_client._client("npx -y server /tmp", None)
+
+
+def test_client_builds_stdio_when_enabled_without_spawning(monkeypatch):
+ _enable(monkeypatch)
+ # Constructing the Client must not spawn the subprocess (spawn happens on
+ # __aenter__); we only assert it builds.
+ client = mcp_client._client("npx -y server /tmp", {"K": "v"})
+ assert client is not None
+
+
+def test_client_http_unaffected_by_gate(monkeypatch):
+ _disable(monkeypatch)
+ assert mcp_client._client("https://example.com/mcp", None) is not None
+
+
+# ── P3: OAuth normalised off for stdio (create + update) ────────────
+
+
+def test_create_forces_oauth_off_for_stdio(tmp_path, monkeypatch):
+ import routes.mcp_servers as routes_mcp
+ from models.mcp_servers import McpServerCreate
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ resp = asyncio.run(
+ routes_mcp.create_mcp_server(
+ McpServerCreate(
+ display_name = "FS", url = "npx -y server /tmp", use_oauth = True
+ ),
+ current_subject = "u",
+ )
+ )
+ assert resp.use_oauth is False
+ assert mcp_servers_db.get_server(resp.id)["use_oauth"] == 0
+
+
+def test_create_keeps_oauth_for_http(tmp_path, monkeypatch):
+ import routes.mcp_servers as routes_mcp
+ from models.mcp_servers import McpServerCreate
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ resp = asyncio.run(
+ routes_mcp.create_mcp_server(
+ McpServerCreate(display_name = "GH", url = "https://gh/mcp", use_oauth = True),
+ current_subject = "u",
+ )
+ )
+ assert resp.use_oauth is True
+
+
+def test_update_url_to_stdio_clears_oauth(tmp_path, monkeypatch):
+ import routes.mcp_servers as routes_mcp
+ from models.mcp_servers import McpServerUpdate
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
+ monkeypatch.setattr(
+ routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0)
+ )
+ mcp_servers_db.create_server(
+ id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True
+ )
+ resp = asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1", McpServerUpdate(url = "npx -y server /tmp"), current_subject = "u"
+ )
+ )
+ assert resp.use_oauth is False
+
+
+# ── P4: env/headers dropped on a transport-type switch ──────────────
+
+
+def test_switch_stdio_to_http_drops_env(tmp_path, monkeypatch):
+ import routes.mcp_servers as routes_mcp
+ from models.mcp_servers import McpServerUpdate
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ mcp_servers_db.create_server(
+ id = "s1",
+ display_name = "A",
+ url = "npx server",
+ headers_json = '{"API_KEY": "secret"}',
+ )
+ resp = asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1", McpServerUpdate(url = "https://remote/mcp"), current_subject = "u"
+ )
+ )
+ # the stdio env must NOT survive as HTTP headers on the remote endpoint
+ assert resp.headers == {}
+ assert mcp_servers_db.get_server("s1")["headers_json"] is None
+
+
+def test_switch_keeps_explicitly_supplied_headers(tmp_path, monkeypatch):
+ import routes.mcp_servers as routes_mcp
+ from models.mcp_servers import McpServerUpdate
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ mcp_servers_db.create_server(
+ id = "s1",
+ display_name = "A",
+ url = "npx server",
+ headers_json = '{"API_KEY": "secret"}',
+ )
+ resp = asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1",
+ McpServerUpdate(
+ url = "https://remote/mcp", headers = {"Authorization": "Bearer new"}
+ ),
+ current_subject = "u",
+ )
+ )
+ assert resp.headers == {"Authorization": "Bearer new"}
+
+
+def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch):
+ import routes.mcp_servers as routes_mcp
+ from models.mcp_servers import McpServerUpdate
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ mcp_servers_db.create_server(
+ id = "s1",
+ display_name = "A",
+ url = "npx server",
+ headers_json = '{"API_KEY": "secret"}',
+ )
+ # editing only the display name (still stdio) must not wipe env vars
+ resp = asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1", McpServerUpdate(display_name = "B"), current_subject = "u"
+ )
+ )
+ assert resp.headers == {"API_KEY": "secret"}
+
+
+# ── P5: reject a command whose first token is a URL scheme ───────────
+
+
+def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch):
+ from routes.mcp_servers import _validate_url
+
+ _enable(monkeypatch)
+ for bad in ["ftp://host/x", "file:///etc/passwd", "ws://h/y"]:
+ with pytest.raises(HTTPException) as exc:
+ _validate_url(bad)
+ assert exc.value.status_code == 400
+
+
+def test_validate_url_allows_url_in_argument(monkeypatch):
+ from routes.mcp_servers import _validate_url
+
+ _enable(monkeypatch)
+ # :// inside an ARGUMENT (not the first token) is still a valid command
+ assert _validate_url("npx server --url https://x/mcp") == (
+ "npx server --url https://x/mcp"
+ )
+
+
+# ── P6: Data Recipe stdio path obeys the same host gate ─────────────
+# build_mcp_providers needs the data_designer plugin, which is only installed in
+# the Studio test job; skip there rather than fail the core matrix.
+
+_STDIO_RECIPE = {
+ "mcp_providers": [
+ {
+ "provider_type": "stdio",
+ "name": "fs",
+ "command": "npx",
+ "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
+ "env": {},
+ }
+ ]
+}
+
+
+def test_data_recipe_skips_stdio_when_disabled(monkeypatch):
+ pytest.importorskip("data_designer")
+ _disable(monkeypatch)
+ from core.data_recipe.service import build_mcp_providers
+
+ # gate off -> the stdio provider is dropped (no subprocess can be spawned)
+ assert build_mcp_providers(_STDIO_RECIPE) == []
+
+
+def test_data_recipe_builds_stdio_when_enabled(monkeypatch):
+ pytest.importorskip("data_designer")
+ _enable(monkeypatch)
+ from core.data_recipe.service import build_mcp_providers
+
+ built = build_mcp_providers(_STDIO_RECIPE)
+ assert len(built) == 1 # constructed (not spawned) only when enabled
diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py
new file mode 100644
index 0000000000..d9a9510378
--- /dev/null
+++ b/studio/backend/tests/test_mcp_stdio_pr5863.py
@@ -0,0 +1,367 @@
+"""Verification tests for PR #5863 (stdio MCP server support).
+
+Covers the pure helpers (is_stdio / parse_stdio_command / stdio_mcp_enabled /
+probe_timeout), the route-level _validate_url gate, and - most importantly -
+that the UNSLOTH_STUDIO_ALLOW_STDIO_MCP gate blocks the stdio transport at all
+five enforcement points (create, update, test, refresh, discovery, execute)
+when disabled, and reaches it when enabled. The transport (_client) is stubbed
+so no real subprocess is spawned; a recorder asserts whether it was reached.
+
+Run from studio/backend: python -m pytest tests/test_mcp_stdio_pr5863.py -q
+"""
+
+import sys
+
+import pytest
+from fastapi import HTTPException
+
+from core.inference import mcp_client
+from storage import mcp_servers_db
+
+
+def _reset_db(tmp_path, monkeypatch):
+ monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
+ monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
+
+
+def _enable(monkeypatch):
+ monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
+
+
+def _disable(monkeypatch):
+ monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
+
+
+# ── transport stub + recorder ───────────────────────────────────────
+
+
+class _FakeTool:
+ def __init__(self, name):
+ self._name = name
+
+ def model_dump(self, exclude_none = True):
+ return {"name": self._name, "description": f"{self._name} tool"}
+
+
+class _Block:
+ def __init__(self, text):
+ self.type = "text"
+ self.text = text
+
+
+class _FakeResult:
+ is_error = False
+
+ def __init__(self, text):
+ self.content = [_Block(text)]
+
+
+class _RecordingClient:
+ """Stands in for fastmcp.Client; records that the transport was opened."""
+
+ def __init__(self, url, headers, use_oauth, recorder):
+ recorder.append({"url": url, "headers": headers, "use_oauth": use_oauth})
+
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, *a):
+ return False
+
+ async def list_tools(self):
+ return [_FakeTool("list_directory"), _FakeTool("write_file")]
+
+ async def call_tool(self, name, args):
+ return _FakeResult(f"called {name}")
+
+
+@pytest.fixture
+def transport(monkeypatch):
+ """Patch mcp_client._client with a recorder. Returns the recorder list;
+ empty == the stdio transport was never reached."""
+ recorder = []
+ monkeypatch.setattr(
+ mcp_client,
+ "_client",
+ lambda url, headers, use_oauth = False: _RecordingClient(
+ url, headers, use_oauth, recorder
+ ),
+ )
+ return recorder
+
+
+# ── 1. is_stdio ─────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+ "addr",
+ [
+ "http://localhost:8000/mcp",
+ "https://example.com/mcp",
+ " https://example.com/mcp ",
+ "HTTPS://EXAMPLE.COM/mcp",
+ ],
+)
+def test_is_stdio_false_for_http(addr):
+ assert mcp_client.is_stdio(addr) is False
+
+
+@pytest.mark.parametrize(
+ "addr",
+ [
+ "npx -y @modelcontextprotocol/server-filesystem /tmp",
+ "python -m some.module",
+ "uvx some-server --flag",
+ "/usr/local/bin/my-server",
+ ],
+)
+def test_is_stdio_true_for_commands(addr):
+ assert mcp_client.is_stdio(addr) is True
+
+
+# ── 2. parse_stdio_command ──────────────────────────────────────────
+
+
+def test_parse_basic_argv():
+ assert mcp_client.parse_stdio_command(
+ "npx -y @modelcontextprotocol/server-filesystem /tmp"
+ ) == ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
+
+
+def test_parse_keeps_url_argument_as_one_command():
+ # gemini "high": a :// inside an ARGUMENT must not break the command.
+ assert mcp_client.parse_stdio_command(
+ "npx server --endpoint https://example.com/mcp"
+ ) == ["npx", "server", "--endpoint", "https://example.com/mcp"]
+
+
+def test_parse_quoted_arg():
+ assert mcp_client.parse_stdio_command('python -m mod --name "a b"') == [
+ "python",
+ "-m",
+ "mod",
+ "--name",
+ "a b",
+ ]
+
+
+def test_parse_empty_returns_empty_list():
+ assert mcp_client.parse_stdio_command(" ") == []
+
+
+def test_parse_unclosed_quote_raises_valueerror():
+ with pytest.raises(ValueError):
+ mcp_client.parse_stdio_command('npx "unclosed')
+
+
+def test_parse_windows_strips_wrapping_quotes(monkeypatch):
+ # gemini "medium": posix=False keeps backslash paths but also the wrapping
+ # quotes; the PR strips a matched pair so argv[0] reaches the OS clean.
+ monkeypatch.setattr(sys, "platform", "win32")
+ parts = mcp_client.parse_stdio_command(
+ r'"C:\Program Files\node\node.exe" server.js'
+ )
+ assert parts[0] == r"C:\Program Files\node\node.exe"
+ assert parts[1] == "server.js"
+
+
+# ── 3. stdio_mcp_enabled ────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("val", ["0", "false", "true", "", " 1 ", "yes", "2"])
+def test_stdio_disabled_for_non_exact_one(monkeypatch, val):
+ monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", val)
+ assert mcp_client.stdio_mcp_enabled() is False
+
+
+def test_stdio_enabled_only_for_exact_one(monkeypatch):
+ _disable(monkeypatch)
+ assert mcp_client.stdio_mcp_enabled() is False
+ monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
+ assert mcp_client.stdio_mcp_enabled() is True
+
+
+# ── 4. probe_timeout ────────────────────────────────────────────────
+
+
+def test_probe_timeout_matrix():
+ assert mcp_client.probe_timeout("https://x/mcp", False) == 8.0
+ assert mcp_client.probe_timeout("https://x/mcp", True) == 305.0
+ assert mcp_client.probe_timeout("npx server", False) == 60.0
+ # oauth wins regardless of address kind (documented behaviour)
+ assert mcp_client.probe_timeout("npx server", True) == 305.0
+
+
+# ── 5. _validate_url gate ───────────────────────────────────────────
+
+
+def test_validate_url_gate_off_rejects_stdio(monkeypatch):
+ _disable(monkeypatch)
+ from routes.mcp_servers import _validate_url
+
+ assert _validate_url("https://example.com/mcp") == "https://example.com/mcp"
+ for bad in ["npx server", "python -m mod", "ftp://host"]:
+ with pytest.raises(HTTPException) as exc:
+ _validate_url(bad)
+ assert exc.value.status_code == 400
+
+
+def test_validate_url_gate_on_accepts_stdio(monkeypatch):
+ _enable(monkeypatch)
+ from routes.mcp_servers import _validate_url
+
+ assert _validate_url("npx -y server /tmp") == "npx -y server /tmp"
+ # http still works when stdio is on
+ assert _validate_url("https://x/mcp") == "https://x/mcp"
+ # url-bearing argument accepted as a command
+ assert _validate_url("npx server --url https://x/mcp") == (
+ "npx server --url https://x/mcp"
+ )
+ # empty / unparseable still rejected
+ for bad in [" ", '"unclosed']:
+ with pytest.raises(HTTPException) as exc:
+ _validate_url(bad)
+ assert exc.value.status_code == 400
+
+
+# ── 6. gate enforcement at every spawn path (mocked transport) ──────
+
+
+def test_create_route_gate(tmp_path, monkeypatch, transport):
+ import asyncio
+
+ from models.mcp_servers import McpServerCreate
+ import routes.mcp_servers as routes_mcp
+
+ _reset_db(tmp_path, monkeypatch)
+ payload = McpServerCreate(display_name = "FS", url = "npx -y server /tmp")
+
+ _disable(monkeypatch)
+ with pytest.raises(HTTPException) as exc:
+ asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u"))
+ assert exc.value.status_code == 400
+
+ _enable(monkeypatch)
+ resp = asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u"))
+ assert resp.url == "npx -y server /tmp"
+
+
+def test_update_http_to_stdio_blocked_when_off(tmp_path, monkeypatch):
+ import asyncio
+
+ from models.mcp_servers import McpServerUpdate
+ import routes.mcp_servers as routes_mcp
+
+ _reset_db(tmp_path, monkeypatch)
+ _disable(monkeypatch)
+ mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp")
+ # editing url -> stdio command must 400 (http->stdio edit bypass closed)
+ with pytest.raises(HTTPException) as exc:
+ asyncio.run(
+ routes_mcp.update_mcp_server(
+ "s1", McpServerUpdate(url = "npx server"), current_subject = "u"
+ )
+ )
+ assert exc.value.status_code == 400
+
+
+def test_test_route_gate(tmp_path, monkeypatch, transport):
+ import asyncio
+
+ from models.mcp_servers import McpServerTestRequest
+ import routes.mcp_servers as routes_mcp
+
+ _reset_db(tmp_path, monkeypatch)
+ req = McpServerTestRequest(url = "npx -y server /tmp")
+
+ _disable(monkeypatch)
+ with pytest.raises(HTTPException) as exc:
+ asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u"))
+ assert exc.value.status_code == 400
+ assert transport == [] # transport never opened
+
+ _enable(monkeypatch)
+ res = asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u"))
+ assert res.ok and res.tool_count == 2
+ assert len(transport) == 1
+
+
+def test_refresh_route_gate(tmp_path, monkeypatch, transport):
+ import asyncio
+
+ import routes.mcp_servers as routes_mcp
+
+ _reset_db(tmp_path, monkeypatch)
+ # a stdio row as if carried over from a desktop DB
+ mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server")
+
+ _disable(monkeypatch)
+ with pytest.raises(HTTPException) as exc:
+ asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u"))
+ assert exc.value.status_code == 400
+ assert transport == []
+
+ _enable(monkeypatch)
+ res = asyncio.run(
+ routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")
+ )
+ assert res.ok and res.tool_count == 2
+ assert len(transport) == 1
+
+
+def test_discovery_gate(tmp_path, monkeypatch, transport):
+ import asyncio
+
+ from core.inference.tools import get_enabled_mcp_tools
+
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(
+ id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True
+ )
+
+ _disable(monkeypatch)
+ assert asyncio.run(get_enabled_mcp_tools()) == []
+ assert transport == [] # filtered out before any probe
+
+ _enable(monkeypatch)
+ specs = asyncio.run(get_enabled_mcp_tools())
+ assert len(specs) == 2
+ assert len(transport) == 1
+
+
+def test_execute_gate(tmp_path, monkeypatch, transport):
+ from core.inference.tools import execute_tool
+
+ _reset_db(tmp_path, monkeypatch)
+ mcp_servers_db.create_server(
+ id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True
+ )
+
+ _disable(monkeypatch)
+ out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"})
+ assert "disabled on this host" in out
+ assert transport == []
+
+ _enable(monkeypatch)
+ out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"})
+ assert out == "called list_directory"
+ assert len(transport) == 1
+
+
+# ── 7. env vars ride headers_json as the subprocess env ─────────────
+
+
+def test_stdio_env_passed_through(tmp_path, monkeypatch, transport):
+ from core.inference.tools import execute_tool
+
+ _reset_db(tmp_path, monkeypatch)
+ _enable(monkeypatch)
+ mcp_servers_db.create_server(
+ id = "stdio1",
+ display_name = "FS",
+ url = "npx server",
+ headers_json = '{"API_KEY": "sk-test"}',
+ is_enabled = True,
+ )
+ execute_tool("mcp__stdio1__list_directory", {})
+ assert transport[-1]["headers"] == {"API_KEY": "sk-test"}
diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py
index 98c7bdaa55..c36363b1ae 100644
--- a/studio/backend/tests/test_mlx_training_worker_config.py
+++ b/studio/backend/tests/test_mlx_training_worker_config.py
@@ -66,6 +66,7 @@ def _load_worker_module():
_worker = _load_worker_module()
_normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer
_normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler
+_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size
def test_mlx_studio_optimizer_aliases_are_explicit():
@@ -82,3 +83,14 @@ def test_mlx_studio_rejects_unknown_optimizer():
def test_mlx_studio_rejects_unknown_scheduler():
with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"):
_normalize_mlx_studio_scheduler("linear_typo")
+
+
+def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer():
+ assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256)
+ assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512)
+ assert _mlx_vlm_max_resized_size(1000, 1000, 512) == (512, 512)
+ assert _mlx_vlm_max_resized_size(256, 128, 1536) == (256, 128)
+ assert _mlx_vlm_max_resized_size(512, 256, 512) == (512, 256)
+ # Half-pixel cases must match the Torch collator (not banker's round).
+ assert _mlx_vlm_max_resized_size(333, 1000, 500) == (167, 500)
+ assert _mlx_vlm_max_resized_size(1000, 333, 500) == (500, 167)
diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py
index 3d179371e3..0b1a65c69e 100644
--- a/studio/backend/tests/test_openai_code_execution.py
+++ b/studio/backend/tests/test_openai_code_execution.py
@@ -269,7 +269,15 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch):
assert len(ends) == 1
assert starts[0]["tool_name"] == "code_execution"
assert starts[0]["tool_call_id"] == "scall_1"
- assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la"}
+ # `_server_tool: True` is the synthetic-builtin marker the
+ # backend stamps onto every provider-side tool_start so the
+ # frontend serializer can distinguish hosted tools from
+ # user-declared functions on history replay.
+ assert starts[0]["arguments"] == {
+ "kind": "bash",
+ "command": "ls -la",
+ "_server_tool": True,
+ }
assert ends[0]["tool_call_id"] == "scall_1"
assert "total 24" in ends[0]["result"]
diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py
index 6d415316d4..f5dee2561a 100644
--- a/studio/backend/tests/test_openai_image_generation.py
+++ b/studio/backend/tests/test_openai_image_generation.py
@@ -207,9 +207,12 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch):
ends = [e for e in image_events if e.get("type") == "tool_end"]
assert len(starts) == 1, image_events
assert len(ends) == 1, image_events
+ # `_server_tool: True` marks this as a provider-side synthetic
+ # tool card on the frontend's history serializer.
assert starts[0]["arguments"] == {
"kind": "image",
"prompt": "A photorealistic cat sitting",
+ "_server_tool": True,
"openai_image_generation_call_id": "img_abc",
}
assert ends[0]["image_b64"] == "AAAA"
diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py
index 22ccba7058..f177ed5ef3 100644
--- a/studio/backend/tests/test_openai_responses_translation.py
+++ b/studio/backend/tests/test_openai_responses_translation.py
@@ -215,6 +215,254 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch):
assert payloads[-1] == "[DONE]"
+def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypatch):
+ """Round 12: caller-supplied function tools forwarded into /v1/responses
+ must have their `function_call` output items translated back into Chat
+ Completions delta.tool_calls, and the terminal chunk must emit
+ finish_reason="tool_calls" so the frontend's accumulator runs the
+ function instead of seeing finish_reason="stop"."""
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ events = [
+ {"type": "response.created"},
+ {
+ "type": "response.output_item.done",
+ "item": {
+ "type": "function_call",
+ "id": "fc_abc",
+ "call_id": "call_xyz",
+ "name": "get_weather",
+ "arguments": '{"city":"SF"}',
+ },
+ },
+ {"type": "response.completed", "response": {}},
+ ]
+ return httpx.Response(
+ 200,
+ content = _responses_sse(events),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http_client(monkeypatch, handler)
+
+ async def run():
+ client = _make_client()
+ lines = await _collect(
+ client._stream_openai_responses(
+ messages = [{"role": "user", "content": "weather?"}],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = None,
+ enable_thinking = None,
+ reasoning_effort = None,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "parameters": {
+ "type": "object",
+ "properties": {"city": {"type": "string"}},
+ },
+ },
+ }
+ ],
+ )
+ )
+ await client.close()
+ return lines
+
+ lines = _drive(run())
+ payloads = [
+ json.loads(line[len("data:") :].strip())
+ for line in lines
+ if line.startswith("data:") and line[len("data:") :].strip() != "[DONE]"
+ ]
+ tool_call_deltas = [
+ p
+ for p in payloads
+ if isinstance(p, dict)
+ and p.get("choices")
+ and p["choices"][0].get("delta", {}).get("tool_calls")
+ ]
+ assert tool_call_deltas, payloads
+ tc = tool_call_deltas[0]["choices"][0]["delta"]["tool_calls"][0]
+ assert tc["id"] == "call_xyz"
+ assert tc["function"]["name"] == "get_weather"
+ assert tc["function"]["arguments"] == '{"city":"SF"}'
+ # Final chunk reports tool_calls instead of stop.
+ terminal = next(
+ p
+ for p in payloads
+ if isinstance(p, dict)
+ and p.get("choices")
+ and p["choices"][0].get("finish_reason") in ("stop", "tool_calls")
+ )
+ assert terminal["choices"][0]["finish_reason"] == "tool_calls", payloads
+
+
+def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch):
+ """Round 13: parallel function_call items must land on distinct
+ delta.tool_calls[].index slots so index-keyed clients don't
+ collapse the second call into the first."""
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ events = [
+ {"type": "response.created"},
+ {
+ "type": "response.output_item.done",
+ "item": {
+ "type": "function_call",
+ "id": "fc_a",
+ "call_id": "call_a",
+ "name": "lookup_a",
+ "arguments": "{}",
+ },
+ },
+ {
+ "type": "response.output_item.done",
+ "item": {
+ "type": "function_call",
+ "id": "fc_b",
+ "call_id": "call_b",
+ "name": "lookup_b",
+ "arguments": "{}",
+ },
+ },
+ {"type": "response.completed", "response": {}},
+ ]
+ return httpx.Response(
+ 200,
+ content = _responses_sse(events),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http_client(monkeypatch, handler)
+
+ async def run():
+ client = _make_client()
+ lines = await _collect(
+ client._stream_openai_responses(
+ messages = [{"role": "user", "content": "x"}],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = None,
+ enable_thinking = None,
+ reasoning_effort = None,
+ tools = [
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup_a",
+ "parameters": {"type": "object"},
+ },
+ },
+ {
+ "type": "function",
+ "function": {
+ "name": "lookup_b",
+ "parameters": {"type": "object"},
+ },
+ },
+ ],
+ )
+ )
+ await client.close()
+ return lines
+
+ lines = _drive(run())
+ indices: list[int] = []
+ for raw in lines:
+ if not raw.startswith("data:"):
+ continue
+ payload = raw[len("data:") :].strip()
+ if payload == "[DONE]":
+ continue
+ try:
+ obj = json.loads(payload)
+ except Exception:
+ continue
+ delta = (obj.get("choices") or [{}])[0].get("delta") or {}
+ for tc in delta.get("tool_calls") or []:
+ indices.append(tc.get("index"))
+ assert indices == [0, 1], indices
+
+
+def test_responses_follow_up_tool_result_uses_function_call_output_items(monkeypatch):
+ """Round 13: a second turn after a Responses function call must
+ serialize the tool_calls history and tool result as Responses
+ `function_call` / `function_call_output` input items, not as
+ Chat Completions role="tool" content."""
+ captured: dict = {}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content.decode("utf-8"))
+ return httpx.Response(
+ 200,
+ content = _responses_sse(
+ [
+ {"type": "response.created"},
+ {"type": "response.completed", "response": {}},
+ ]
+ ),
+ headers = {"content-type": "text/event-stream"},
+ )
+
+ _mock_http_client(monkeypatch, handler)
+
+ async def run():
+ client = _make_client()
+ await _collect(
+ client._stream_openai_responses(
+ messages = [
+ {"role": "user", "content": "weather?"},
+ {
+ "role": "assistant",
+ "content": "",
+ "tool_calls": [
+ {
+ "id": "call_xyz",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city":"SF"}',
+ },
+ }
+ ],
+ },
+ {
+ "role": "tool",
+ "tool_call_id": "call_xyz",
+ "content": "sunny",
+ },
+ {"role": "user", "content": "thanks"},
+ ],
+ model = "gpt-5.5",
+ temperature = 0.7,
+ top_p = 0.95,
+ max_tokens = None,
+ enable_thinking = None,
+ reasoning_effort = None,
+ )
+ )
+ await client.close()
+
+ _drive(run())
+ items = captured["body"]["input"]
+ types = [it.get("type") or it.get("role") for it in items]
+ assert "function_call" in types, items
+ assert "function_call_output" in types, items
+ fc = next(it for it in items if it.get("type") == "function_call")
+ assert fc["call_id"] == "call_xyz"
+ assert fc["name"] == "get_weather"
+ assert fc["arguments"] == '{"city":"SF"}'
+ fco = next(it for it in items if it.get("type") == "function_call_output")
+ assert fco["call_id"] == "call_xyz"
+ assert fco["output"] == "sunny"
+
+
def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
events = [
diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py
new file mode 100644
index 0000000000..2ce9b55789
--- /dev/null
+++ b/studio/backend/tests/test_rocm_oom_guard.py
@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier).
+
+Covers the three classification paths:
+ Path 1 – canonical gcnArchName attribute present.
+ Path 2 – gcnArchName absent, alternate-spelling attribute present.
+ Path 3 – ALL arch attrs absent; falls back to device-name substring match.
+
+Regression for: Strix Halo (gfx1151) misclassified as discrete on AMD SDK /
+Radeon wheels that populate props.name = "Radeon 8060S Graphics" but do NOT
+set any gcnArchName attribute. Without the 8060s/8050s name patterns the
+fallback returned is_unified=False, applying the 0.90 fraction instead of
+0.80 and leaving only ~12.8 GiB OS headroom on a 128 GiB unified-memory pool.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from core.training.worker import _rocm_classify_unified_memory
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+
+def _props(**kwargs) -> SimpleNamespace:
+ """Build a fake device-properties object with the given attributes."""
+ return SimpleNamespace(**kwargs)
+
+
+# ── Path 1: canonical gcnArchName ────────────────────────────────────────────
+
+
+class TestCanonicalGcnArchName:
+ """gcnArchName is present and populated."""
+
+ @pytest.mark.parametrize(
+ "arch, expected_unified",
+ [
+ ("gfx1150", True), # Strix Point
+ ("gfx1151", True), # Strix Halo
+ ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete
+ ("gfx906", False), # MI50 — discrete server GPU
+ ("gfx1201", False), # RX 9070 XT — discrete
+ ],
+ )
+ def test_canonical_attr(self, arch: str, expected_unified: bool) -> None:
+ props = _props(gcnArchName = arch, name = "irrelevant")
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == arch
+ assert is_unified is expected_unified
+
+ def test_arch_with_colon_suffix_stripped(self) -> None:
+ """gcnArchName can carry xnack/sramecc suffix; only the base is kept."""
+ props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant")
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == "gfx1151"
+ assert is_unified is True
+
+ def test_canonical_attr_wins_over_name(self) -> None:
+ """Arch attr takes priority; device name should be ignored."""
+ # Discrete arch, but name looks like a unified SKU — arch must win.
+ props = _props(gcnArchName = "gfx1100", name = "Radeon 890M")
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == "gfx1100"
+ assert is_unified is False
+
+
+# ── Path 2: alternate-spelling fallback ──────────────────────────────────────
+
+
+class TestAlternateSpellingFallback:
+ """gcnArchName is missing but an alternate attr spelling is present."""
+
+ @pytest.mark.parametrize(
+ "attr_name",
+ ["gcn_arch_name", "arch_name", "gfx_arch_name"],
+ )
+ def test_alternate_attr_unified(self, attr_name: str) -> None:
+ props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics")
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == "gfx1151"
+ assert is_unified is True
+
+ @pytest.mark.parametrize(
+ "attr_name",
+ ["gcn_arch_name", "arch_name", "gfx_arch_name"],
+ )
+ def test_alternate_attr_discrete(self, attr_name: str) -> None:
+ props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT")
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == "gfx1201"
+ assert is_unified is False
+
+ def test_first_non_empty_attr_wins(self) -> None:
+ """When multiple alternate attrs are present the first non-empty one wins."""
+ props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant")
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == "gfx1151"
+ assert is_unified is True
+
+
+# ── Path 3: device-name fallback ─────────────────────────────────────────────
+
+
+class TestDeviceNameFallback:
+ """ALL arch attrs absent — classifier must rely solely on device name."""
+
+ # --- unified-memory devices that MUST be detected ---
+
+ @pytest.mark.parametrize(
+ "device_name",
+ [
+ # gfx1150 Strix Point
+ "Radeon 890M",
+ "AMD Radeon 890M Graphics",
+ "RADEON 890M", # case-insensitive
+ "Radeon 880M",
+ "AMD Radeon 880M Graphics",
+ # gfx1151 Strix Halo — the regression case from the review
+ "Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch)
+ "AMD Radeon 8060S",
+ "Radeon 8050S Graphics", # cut-down Strix Halo SKU
+ "AMD Radeon 8050S",
+ # case variants
+ "RADEON 8060S GRAPHICS",
+ "radeon 8050s",
+ ],
+ )
+ def test_unified_memory_detected(self, device_name: str) -> None:
+ props = _props(name = device_name)
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == "", f"expected empty gcn_arch, got {gcn!r}"
+ assert (
+ is_unified is True
+ ), f"device {device_name!r} should be classified as unified-memory"
+
+ # --- discrete devices that must NOT be mis-classified ---
+
+ @pytest.mark.parametrize(
+ "device_name",
+ [
+ "Radeon RX 9070 XT",
+ "AMD Radeon RX 7900 XTX",
+ "Radeon RX 6900 XT",
+ "Radeon Pro W7900",
+ "AMD Instinct MI300X",
+ # Names that contain superficially similar substrings but are discrete
+ "Radeon RX 580",
+ "Radeon VII",
+ ],
+ )
+ def test_discrete_not_misclassified(self, device_name: str) -> None:
+ props = _props(name = device_name)
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == ""
+ assert (
+ is_unified is False
+ ), f"discrete device {device_name!r} should NOT be classified as unified-memory"
+
+ def test_empty_name_returns_false(self) -> None:
+ """Completely absent name must not crash and must default to discrete."""
+ props = _props() # no 'name' attr at all
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == ""
+ assert is_unified is False
+
+ def test_none_name_returns_false(self) -> None:
+ props = _props(name = None)
+ gcn, is_unified = _rocm_classify_unified_memory(props)
+ assert gcn == ""
+ assert is_unified is False
diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py
index 7ffa9bb384..6df491610b 100644
--- a/studio/backend/tests/test_studio_train_validation.py
+++ b/studio/backend/tests/test_studio_train_validation.py
@@ -18,6 +18,8 @@ from models.training import (
_MAX_LORA_ALPHA,
_MAX_LORA_R,
_MAX_SEQ_LENGTH,
+ _MAX_VISION_IMAGE_SIZE,
+ _MIN_VISION_IMAGE_SIZE,
)
@@ -62,6 +64,52 @@ class TestBatchSizeCap:
_check_field("batch_size", 0)
+class TestVisionImageSizeCap:
+ def test_none_accepts_model_default(self):
+ _check_field("vision_image_size", None)
+
+ @pytest.mark.parametrize(
+ "value",
+ [_MIN_VISION_IMAGE_SIZE, 640, 1000, _MAX_VISION_IMAGE_SIZE],
+ )
+ def test_in_range_accepts(self, value):
+ _check_field("vision_image_size", value)
+ assert _MIN_VISION_IMAGE_SIZE == 256
+ assert _MAX_VISION_IMAGE_SIZE == 2048
+
+ @pytest.mark.parametrize(
+ "value",
+ [_MIN_VISION_IMAGE_SIZE - 1, _MAX_VISION_IMAGE_SIZE + 1, 640.5, True],
+ )
+ def test_invalid_rejects(self, value):
+ with pytest.raises(ValidationError):
+ _check_field("vision_image_size", value)
+
+ @pytest.mark.parametrize("value", [True, False])
+ def test_bool_error_says_integer_not_range(self, value):
+ # Regression guard: bools must say "integer or null", not "in [256, 2048]".
+ with pytest.raises(ValidationError) as exc:
+ _check_field("vision_image_size", value)
+ assert "integer or null" in str(exc.value)
+
+ @pytest.mark.parametrize("value", ["++512", "--256", "+-+512", "+", "-"])
+ def test_multi_sign_string_says_integer_not_raw(self, value):
+ # Regression guard: multi-sign strings must not leak int()'s raw
+ # "invalid literal" message; precise contract is "integer or null".
+ with pytest.raises(ValidationError) as exc:
+ _check_field("vision_image_size", value)
+ assert "integer or null" in str(exc.value)
+ assert "invalid literal" not in str(exc.value)
+
+ @pytest.mark.parametrize("value", ["５１２", "٥١٢", "१०२४"])
+ def test_unicode_digit_string_rejected(self, value):
+ # Full-width / Arabic-Indic / Devanagari digits must be rejected so the
+ # value reaching the backend equals the ASCII the user typed.
+ with pytest.raises(ValidationError) as exc:
+ _check_field("vision_image_size", value)
+ assert "integer or null" in str(exc.value)
+
+
class TestLoraRCap:
def test_at_cap_accepts(self):
_check_field("lora_r", _MAX_LORA_R)
diff --git a/studio/backend/utils/cpu_threads.py b/studio/backend/utils/cpu_threads.py
new file mode 100644
index 0000000000..922f95ce55
--- /dev/null
+++ b/studio/backend/utils/cpu_threads.py
@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Early CPU thread-pool configuration for Studio processes."""
+
+import os
+from typing import MutableMapping, Optional
+
+
+_THREAD_POOL_ENV_VARS = (
+ "OMP_NUM_THREADS",
+ "MKL_NUM_THREADS",
+ "OPENBLAS_NUM_THREADS",
+ "NUMEXPR_NUM_THREADS",
+)
+
+
+def configure_cpu_threads(env: Optional[MutableMapping[str, str]] = None) -> None:
+ """Apply ``UNSLOTH_CPU_THREADS`` to native CPU pools when configured.
+
+ This must run before importing libraries that initialize an OpenMP or
+ BLAS thread pool. Library-specific variables are left untouched so users
+ can override a single runtime independently.
+ """
+ environ = os.environ if env is None else env
+ configured = environ.get("UNSLOTH_CPU_THREADS", "").strip()
+ if not configured:
+ return
+
+ try:
+ thread_count = int(configured)
+ except ValueError as exc:
+ raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer") from exc
+ if thread_count < 1:
+ raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer")
+
+ value = str(thread_count)
+ for variable in _THREAD_POOL_ENV_VARS:
+ environ.setdefault(variable, value)
diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py
index fdb1ab4520..48d5890399 100644
--- a/studio/backend/utils/hardware/amd.py
+++ b/studio/backend/utils/hardware/amd.py
@@ -11,18 +11,35 @@ nvidia.py counterparts.
import json
import math
import os
+import platform
import re
import subprocess
+import sys
from typing import Any, Optional
from loggers import get_logger
from utils.native_path_leases import child_env_without_native_path_secret
+from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = get_logger(__name__)
+# amd-smi on Windows must initialise the full ROCm runtime on first call, which
+# can take 15-25 s on cold hardware. Linux is consistently < 2 s.
+_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10
-def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
+# Circuit breaker: stop calling amd-smi after this many consecutive failures.
+# On Windows, each failed call spawns a process that may show a UAC/DiskPart
+# elevation prompt. Once we know amd-smi doesn't work we stop polling it.
+_AMD_SMI_FAILURE_LIMIT = 3
+_amd_smi_consecutive_failures = 0
+_amd_smi_disabled = False
+
+
+def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]:
"""Run amd-smi with the given arguments and return parsed JSON, or None."""
+ global _amd_smi_consecutive_failures, _amd_smi_disabled
+ if _amd_smi_disabled:
+ return None
try:
result = subprocess.run(
["amd-smi", *args, "--json"],
@@ -30,13 +47,40 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
+ **windows_hidden_subprocess_kwargs(),
)
except (OSError, subprocess.TimeoutExpired) as e:
- logger.warning("amd-smi query failed: %s", e)
+ if isinstance(e, FileNotFoundError):
+ # amd-smi ships with Adrenalin, not the HIP SDK -- absence is
+ # expected on HIP SDK-only Windows setups. Log at debug only.
+ logger.debug("amd-smi not found (not in PATH): %s", e)
+ else:
+ logger.warning("amd-smi query failed: %s", e)
+ _amd_smi_consecutive_failures += 1
+ if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
+ logger.info(
+ "amd-smi not available (not installed; expected on HIP SDK-only systems); "
+ "GPU VRAM polling disabled"
+ )
+ _amd_smi_disabled = True
return None
- if result.returncode != 0 or not result.stdout.strip():
+ if result.returncode != 0:
logger.warning("amd-smi returned code %d", result.returncode)
+ _amd_smi_consecutive_failures += 1
+ if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
+ logger.info(
+ "amd-smi not available (not installed; expected on HIP SDK-only systems); "
+ "GPU VRAM polling disabled"
+ )
+ _amd_smi_disabled = True
return None
+ if not result.stdout.strip():
+ # amd-smi exited successfully but produced no output (e.g. no GPUs
+ # visible on this query, or a version that emits nothing for --json).
+ # This is not a tool failure, so don't count against the circuit breaker.
+ logger.debug("amd-smi exited 0 but returned no output")
+ return None
+ _amd_smi_consecutive_failures = 0 # reset on success
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
@@ -352,7 +396,7 @@ def get_visible_gpu_utilization(
)
parsed_id = _parse_numeric(raw_id)
if parsed_id is None:
- logger.debug(
+ logger.warning(
"amd-smi GPU id %r could not be parsed; falling back to "
"enumeration index %d",
raw_id,
@@ -360,7 +404,15 @@ def get_visible_gpu_utilization(
)
idx = fallback_idx
else:
- idx = int(parsed_id)
+ rounded = round(parsed_id)
+ if rounded != parsed_id:
+ logger.warning(
+ "amd-smi GPU id %r parsed as non-integer %r; truncating to %d",
+ raw_id,
+ parsed_id,
+ rounded,
+ )
+ idx = int(rounded)
if idx not in visible_set:
continue
metrics = _extract_gpu_metrics(gpu_data)
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index ede37e2953..180fde8f13 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -16,8 +16,16 @@ Usage:
...
"""
+import copy
+import gc
+import glob
import os
import platform
+import re
+import subprocess
+import sys
+import types
+from importlib.metadata import PackageNotFoundError, version as pkg_version
import structlog
from loggers import get_logger
from enum import Enum
@@ -120,11 +128,13 @@ def detect_hardware() -> DeviceType:
# Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes.
# DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP.
- if getattr(torch.version, "hip", None) is not None:
+ # AMD's repo.radeon.com SDK wheels (e.g. 2.9.0+rocmsdk20251116) do
+ # not set torch.version.hip, so fall back to checking __version__.
+ _hip_ver = getattr(torch.version, "hip", None)
+ if _hip_ver is not None or "rocm" in torch.__version__.lower():
IS_ROCM = True
- print(
- f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}"
- )
+ _hip_label = _hip_ver or torch.__version__
+ print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}")
else:
print(f"Hardware detected: CUDA -- {device_name}")
return DEVICE
@@ -176,8 +186,6 @@ def clear_gpu_cache():
Clear GPU memory cache for the current device.
Safe to call on any platform — no-ops gracefully.
"""
- import gc
-
gc.collect()
device = get_device()
@@ -359,8 +367,6 @@ def get_package_versions() -> Dict[str, Optional[str]]:
Returns dict with keys: unsloth, torch, transformers, cuda.
Missing packages yield None.
"""
- from importlib.metadata import version as pkg_version, PackageNotFoundError
-
packages = ("unsloth", "torch", "transformers")
versions: Dict[str, Optional[str]] = {}
@@ -466,7 +472,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]:
try:
func = getattr(_backend, func_name)
result = func(*args, **kwargs)
- if result.get("available"):
+ if isinstance(result, dict) and result.get("available"):
return result
except Exception as e:
logger.warning("%s %s query failed: %s", backend_name, func_name, e)
@@ -479,9 +485,6 @@ def _read_apple_gpu_stats() -> Dict[str, Any]:
Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory).
Returns empty dict on failure.
"""
- import subprocess
- import re
-
try:
result = subprocess.run(
["ioreg", "-r", "-c", "AGXAccelerator"],
@@ -506,6 +509,133 @@ def _read_apple_gpu_stats() -> Dict[str, Any]:
}
+def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]:
+ """Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent."""
+ if platform.system() != "Linux":
+ return None
+ try:
+ files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent")
+ if not files:
+ return None
+ values = [int(open(f).read().strip()) for f in files]
+ return round(sum(values) / len(values), 1)
+ except Exception:
+ return None
+
+
+def _rocm_linux_sysfs_temp_c() -> Optional[float]:
+ """Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C)."""
+ if platform.system() != "Linux":
+ return None
+ try:
+ files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input")
+ if not files:
+ return None
+ temps = [int(open(f).read().strip()) / 1000.0 for f in files]
+ return round(max(temps), 1)
+ except Exception:
+ return None
+
+
+def _rocm_linux_sysfs_power_w() -> Optional[float]:
+ """Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts)."""
+ if platform.system() != "Linux":
+ return None
+ try:
+ for pattern in (
+ "/sys/class/drm/card*/device/hwmon/hwmon*/power1_average",
+ "/sys/class/drm/card*/device/hwmon/hwmon*/power1_input",
+ ):
+ files = glob.glob(pattern)
+ if files:
+ watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files)
+ return round(watts, 1)
+ return None
+ except Exception:
+ return None
+
+
+def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]:
+ """Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes)."""
+ if platform.system() != "Windows":
+ return None
+ try:
+ ps = (
+ "$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'"
+ " -ErrorAction SilentlyContinue).CounterSamples;"
+ "if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}"
+ )
+ r = subprocess.run(
+ ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
+ capture_output = True,
+ text = True,
+ timeout = 5,
+ )
+ if r.returncode != 0 or not r.stdout.strip():
+ return None
+ val = float(r.stdout.strip())
+ return round(val, 1) if val >= 0 else None
+ except Exception:
+ return None
+
+
+def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
+ """Query system-wide AMD GPU VRAM via Linux DRM sysfs.
+
+ Reads /sys/class/drm/card*/device/mem_info_vram_* which the kernel
+ updates in real-time across all processes. No tools required.
+ Returns (used_gb, total_gb) or (None, None) on failure.
+ """
+ if platform.system() != "Linux":
+ return None, None
+ try:
+ used_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used")
+ total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total")
+ if not used_files or not total_files:
+ return None, None
+ used_bytes = sum(int(open(f).read().strip()) for f in used_files)
+ total_bytes = sum(int(open(f).read().strip()) for f in total_files)
+ if total_bytes == 0:
+ return None, None
+ return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
+ except Exception:
+ return None, None
+
+
+def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
+ """Query system-wide dedicated GPU VRAM via Windows Performance Counters.
+
+ Uses the same data source as Task Manager so it reflects cross-process
+ usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi.
+ Returns (used_gb, total_gb) or (None, None) on failure.
+ """
+ if platform.system() != "Windows":
+ return None, None
+ try:
+ ps = (
+ "$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
+ " -ErrorAction SilentlyContinue).CounterSamples;"
+ "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}"
+ )
+ r = subprocess.run(
+ ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
+ capture_output = True,
+ text = True,
+ timeout = 5,
+ )
+ if r.returncode != 0 or not r.stdout.strip():
+ return None, None
+ used_bytes = float(r.stdout.strip())
+ if used_bytes < 0:
+ return None, None
+ import torch as _torch
+
+ total_bytes = _torch.cuda.get_device_properties(0).total_memory
+ return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
+ except Exception:
+ return None, None
+
+
def get_gpu_utilization() -> Dict[str, Any]:
"""Return a live snapshot of device utilization information."""
device = get_device()
@@ -514,7 +644,78 @@ def get_gpu_utilization() -> Dict[str, Any]:
result = _smi_query("get_primary_gpu_utilization")
if result is not None:
result["backend"] = _backend_label(device)
+ if IS_ROCM:
+ # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
+ _reconcile_primary_rocm_unified_memory(
+ result, _get_parent_visible_gpu_spec()
+ )
return result
+ # SMI tool unavailable or returned no usable data. On Windows, query
+ # the Performance Counter API (same source as Task Manager) for
+ # system-wide dedicated VRAM — covers cross-process usage that
+ # torch.cuda.mem_get_info cannot see from the Studio server process.
+ if IS_ROCM and platform.system() == "Windows":
+ _win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
+ if _win_used is not None and _win_total is not None:
+ _win_util = _rocm_windows_perf_counter_gpu_util_pct()
+ return {
+ "available": True,
+ "backend": _backend_label(device),
+ "gpu_utilization_pct": _win_util,
+ "temperature_c": None,
+ "vram_used_gb": _win_used,
+ "vram_total_gb": _win_total,
+ "vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
+ if _win_total > 0
+ else None,
+ "power_draw_w": None,
+ "power_limit_w": None,
+ "power_utilization_pct": None,
+ }
+ # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed.
+ if IS_ROCM and platform.system() == "Linux":
+ _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb()
+ if _linux_used is not None and _linux_total is not None:
+ _linux_util = _rocm_linux_sysfs_gpu_busy_pct()
+ _linux_temp = _rocm_linux_sysfs_temp_c()
+ _linux_power = _rocm_linux_sysfs_power_w()
+ return {
+ "available": True,
+ "backend": _backend_label(device),
+ "gpu_utilization_pct": _linux_util,
+ "temperature_c": _linux_temp,
+ "vram_used_gb": _linux_used,
+ "vram_total_gb": _linux_total,
+ "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
+ if _linux_total > 0
+ else None,
+ "power_draw_w": _linux_power,
+ "power_limit_w": None,
+ "power_utilization_pct": None,
+ }
+ # Last resort: torch mem_get_info (process-local).
+ _visible_spec = _get_parent_visible_gpu_spec()
+ _numeric_ids = _visible_spec.get("numeric_ids") or [0]
+ _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0]
+ _torch_devices = _torch_get_per_device_info(_primary_idx)
+ if _torch_devices:
+ _td = _torch_devices[0]
+ _total = _td["total_gb"]
+ _used = _td["used_gb"]
+ return {
+ "available": True,
+ "backend": _backend_label(device),
+ "gpu_utilization_pct": None,
+ "temperature_c": None,
+ "vram_used_gb": _used,
+ "vram_total_gb": _total,
+ "vram_utilization_pct": round((_used / _total) * 100, 1)
+ if _total > 0
+ else None,
+ "power_draw_w": None,
+ "power_limit_w": None,
+ "power_utilization_pct": None,
+ }
# MLX path: single _read_apple_gpu_stats() call carries both VRAM-used
# bytes and GPU utilization %. psutil for unified-memory total is cheap.
@@ -578,6 +779,77 @@ def get_gpu_utilization() -> Dict[str, Any]:
return {"available": False, "backend": _backend_label(device)}
+def _apply_unified_memory_correction(
+ device_metrics: Dict[str, Any], torch_info: Dict[str, Any]
+) -> None:
+ """Per-device reconciliation: when torch reports a larger memory total
+ than amd-smi, overwrite the smi VRAM fields in place.
+
+ Used by both the multi-device and primary-device reconciliation helpers
+ so the two endpoints stay in sync on AMD iGPUs with unified memory.
+ """
+ torch_total_gb = torch_info["total_gb"]
+ smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
+ if torch_total_gb > smi_total_gb:
+ torch_used_gb = torch_info["used_gb"]
+ device_metrics["vram_total_gb"] = torch_total_gb
+ device_metrics["vram_used_gb"] = torch_used_gb
+ device_metrics["vram_utilization_pct"] = (
+ round((torch_used_gb / torch_total_gb) * 100, 1)
+ if torch_total_gb > 0
+ else None
+ )
+ logger.debug(
+ "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with "
+ "torch mem_get_info total (%.2f GB) for device %s",
+ smi_total_gb,
+ torch_total_gb,
+ torch_info.get("index"),
+ )
+
+
+def _reconcile_rocm_unified_memory(
+ utilization: Dict[str, Any], device_indices: list[int]
+) -> None:
+ """Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo).
+
+ amd-smi reports only the dedicated slice (~512 MB); torch sees the full
+ GTT pool (~128 GB). When torch total > smi total, overwrite per-device
+ VRAM fields so GPU selection uses the real available memory.
+ """
+ torch_devices = _torch_get_per_device_info(device_indices)
+ if not torch_devices:
+ return
+ torch_by_index = {td["index"]: td for td in torch_devices}
+ for dev in utilization.get("devices", []):
+ td = torch_by_index.get(dev.get("index"))
+ if td is None:
+ continue
+ _apply_unified_memory_correction(dev, td)
+
+
+def _reconcile_primary_rocm_unified_memory(
+ utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any]
+) -> None:
+ """Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict."""
+ numeric_ids = parent_visible_spec.get("numeric_ids")
+ if numeric_ids is None:
+ # No visibility env var set: torch ordinal 0 is the primary device.
+ primary_idx = [0]
+ elif len(numeric_ids) == 0:
+ # Empty mask (HIP_VISIBLE_DEVICES="" or "-1"): no GPU is visible to
+ # this process. Querying torch device 0 would raise a RuntimeError or
+ # return stale/wrong data, so bail out rather than writing bad values
+ # into the utilization dict.
+ return
+ else:
+ primary_idx = [int(numeric_ids[0])]
+ torch_devices = _torch_get_per_device_info(primary_idx)
+ if not torch_devices:
+ return
+ _apply_unified_memory_correction(utilization, torch_devices[0])
+
+
def get_visible_gpu_utilization() -> Dict[str, Any]:
device = get_device()
@@ -590,6 +862,10 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
)
if result is not None:
result["backend"] = _backend_label(device)
+ numeric_ids = parent_visible_spec.get("numeric_ids")
+ if IS_ROCM and numeric_ids is not None:
+ # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
+ _reconcile_rocm_unified_memory(result, numeric_ids)
return result
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
@@ -689,7 +965,15 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]:
# Use explicit None checks (not `or`) so empty string "" is honoured
# as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES.
cuda_visible = None
- if IS_ROCM:
+ # Prefer ROCm masks only on a ROCm host, or when no CUDA mask is set, so a
+ # stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES.
+ _is_rocm_spec = IS_ROCM or (
+ "CUDA_VISIBLE_DEVICES" not in os.environ
+ and (
+ "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
+ )
+ )
+ if _is_rocm_spec:
hip_vis = os.environ.get("HIP_VISIBLE_DEVICES")
rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES")
if hip_vis is not None:
@@ -865,7 +1149,57 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
def _determine_attention_impl_for_gpu_estimate(config) -> str:
- import copy as _copy
+ # torch.distributed is incomplete on Windows ROCm — torch._C is a C
+ # extension (not a package), so Python cannot import the submodule
+ # torch._C._distributed_c10d that torch.distributed depends on.
+ # Inject an empty stub into sys.modules BEFORE importing torch.distributed
+ # so the import succeeds, then patch the missing process-group helpers.
+ if sys.platform == "win32" and IS_ROCM:
+ # Dummy class for any name torch.distributed tries to import from these stubs
+ class _Dummy:
+ pass
+
+ for _c10d_name in (
+ "torch._C._distributed_c10d",
+ "torch._C._distributed_autograd",
+ "torch._C._distributed_rpc",
+ ):
+ if _c10d_name not in sys.modules:
+ _stub = types.ModuleType(_c10d_name)
+ # torch.distributed imports these names from _distributed_c10d;
+ # provide no-op dummies so the import doesn't raise AttributeError.
+ for _sym in (
+ "FakeProcessGroup",
+ "ProcessGroup",
+ "Work",
+ "Store",
+ "PrefixStore",
+ "FileStore",
+ "TCPStore",
+ "HashStore",
+ "Reducer",
+ "Logger",
+ "DistributedDebugLevel",
+ "GradBucket",
+ "BuiltinCommHookType",
+ ):
+ setattr(_stub, _sym, _Dummy)
+ sys.modules[_c10d_name] = _stub
+
+ try:
+ import torch.distributed as _td
+
+ for _attr, _stub in (
+ ("is_initialized", lambda: False),
+ ("is_available", lambda: False),
+ ("get_rank", lambda: 0),
+ ("get_world_size", lambda: 1),
+ ("is_torchelastic_launched", lambda: False),
+ ):
+ if not hasattr(_td, _attr):
+ setattr(_td, _attr, _stub)
+ except ImportError:
+ pass
from unsloth.models._utils import resolve_attention_implementation
from transformers import AutoModel, AutoModelForCausalLM
@@ -875,7 +1209,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str:
# `sub_configs` and propagates to nested text_config / sub-configs, so a
# shallow copy still mutates those shared inner objects on the cached
# config returned by _load_config_for_gpu_estimate. Deepcopy isolates them.
- config_copy = _copy.deepcopy(config)
+ config_copy = copy.deepcopy(config)
model_class = None
for auto_model in (AutoModelForCausalLM, AutoModel):
@@ -1062,7 +1396,10 @@ def estimate_required_model_memory_gb(
_determine_attention_impl_for_gpu_estimate(config)
)
except Exception as e:
- logger.warning(
+ # Log at debug: on Windows ROCm the torch.distributed stub does
+ # not implement Store, so this fires on every estimate call.
+ # It is expected and non-actionable -- eager is the safe fallback.
+ logger.debug(
"Could not resolve attention implementation for '%s': %s",
estimate_model,
e,
@@ -1552,14 +1889,35 @@ def apply_gpu_ids(gpu_ids) -> None:
# parent process already set a ROCm visibility variable -- that
# way a downstream ROCm process inherits the narrowed mask even
# before Studio's hardware detection has classified the host.
+ # Final fallback: probe torch.version.hip so AMD workers without
+ # HIP_VISIBLE_DEVICES still get the correct ROCm visibility mask.
_inherits_rocm_visibility = (
"HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
)
- if IS_ROCM or _inherits_rocm_visibility:
+ _is_rocm = IS_ROCM or _inherits_rocm_visibility
+ if not _is_rocm:
+ # torch.version.hip is a non-empty string on ROCm, None on CUDA.
+ # AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but
+ # still encode "rocm" in torch.__version__, matching detect_hardware().
+ # Broad except: a probe failure must never crash a training worker.
+ try:
+ import torch as _torch
+
+ _is_rocm = (
+ getattr(_torch.version, "hip", None) is not None
+ or "rocm" in getattr(_torch, "__version__", "").lower()
+ )
+ except Exception as e:
+ logger.debug(
+ "apply_gpu_ids: torch ROCm probe skipped (%s: %s)",
+ type(e).__name__,
+ e,
+ )
+ if _is_rocm:
os.environ["HIP_VISIBLE_DEVICES"] = value
os.environ["ROCR_VISIBLE_DEVICES"] = value
_visible_gpu_count = None
- if IS_ROCM or _inherits_rocm_visibility:
+ if _is_rocm:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)
else:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value)
@@ -1652,8 +2010,6 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
Returns:
A safe integer ≥ 1.
"""
- import sys
-
# Windows and macOS use 'spawn' for multiprocessing -- the overhead of
# re-importing torch/transformers/unsloth per worker is typically slower
# than single-process.
@@ -1704,8 +2060,6 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]:
``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``).
Only ``num_proc=None`` guarantees in-process execution.
"""
- import sys
-
if sys.platform in ("win32", "darwin"):
return None
return safe_num_proc(desired)
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 993995ee57..dc34444ccb 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -1156,13 +1156,18 @@ def detect_gguf_model(path: str) -> Optional[str]:
p = Path(path)
# Case 1: direct .gguf file
- if p.suffix.lower() == ".gguf" and p.is_file():
+ if p.suffix.lower() == ".gguf":
if _is_mmproj(p.name):
return None
- # Use absolute (not resolve) to preserve symlink names -- e.g.
- # Ollama .studio_links/model.gguf -> blobs/sha256-... should
- # keep the readable symlink name, not the opaque blob hash.
- return str(p.absolute())
+ # Extension is authoritative: don't gate on is_file()/exists(), which
+ # can fail in the Windows lock window after llama-server is killed.
+ try:
+ is_dir = p.is_dir()
+ except OSError:
+ is_dir = False # stat() unavailable in the lock window
+ if not is_dir:
+ return str(p.absolute()) # absolute() keeps symlink names readable
+ # Directory named "*.gguf": fall through to the dir scan below.
# Case 2: directory containing .gguf files (skip mmproj)
if p.is_dir():
diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py
index 5c42e890d1..e0ce02261b 100644
--- a/studio/backend/utils/wheel_utils.py
+++ b/studio/backend/utils/wheel_utils.py
@@ -15,6 +15,7 @@ import urllib.request
from typing import Callable
from utils.native_path_leases import child_env_without_native_path_secret
+from utils.subprocess_compat import windows_hidden_subprocess_kwargs
_logger = logging.getLogger(__name__)
@@ -106,6 +107,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
+ **windows_hidden_subprocess_kwargs(),
)
except subprocess.TimeoutExpired:
return None
diff --git a/studio/frontend/package.json b/studio/frontend/package.json
index 83b1fd96f9..b43e174889 100644
--- a/studio/frontend/package.json
+++ b/studio/frontend/package.json
@@ -12,6 +12,7 @@
"lint": "eslint .",
"preview": "vite preview",
"typecheck": "tsc -b --pretty false",
+ "i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts",
"biome:check": "biome check",
"biome:fix": "biome check --write"
},
diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx
index c7bc0440bd..f0a417638d 100644
--- a/studio/frontend/src/app/router.tsx
+++ b/studio/frontend/src/app/router.tsx
@@ -3,6 +3,7 @@
import { Link, createRouter, useRouterState } from "@tanstack/react-router";
import { Button } from "@/components/ui/button";
+import { useT } from "@/i18n";
import { Route as rootRoute } from "./routes/__root";
import { Route as dataRecipesRoute } from "./routes/data-recipes";
import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId";
@@ -31,7 +32,9 @@ const routeTree = rootRoute.addChildren([
]);
function DefaultNotFound() {
+ const t = useT();
const pathname = useRouterState({ select: (s) => s.location.pathname });
+
return (
![]()
);
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index 47bff815e6..57ed233d51 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -6,8 +6,9 @@ import { Navbar } from "@/components/navbar";
import { fetchDeviceType, usePlatformStore } from "@/config/env";
import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
import { SettingsDialog, useSettingsDialogStore } from "@/features/settings";
-import { useTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard";
+import { useTrainingUnloadGuard } from "@/features/training";
import { useSidebarPin } from "@/hooks/use-sidebar-pin";
+import { useT, type TranslationKey } from "@/i18n";
import {
Outlet,
createRootRoute,
@@ -16,24 +17,25 @@ import {
useRouterState,
} from "@tanstack/react-router";
import { AnimatePresence, motion } from "motion/react";
-import { Suspense, useEffect, useLayoutEffect, type ReactNode } from "react";
+import { Suspense, useEffect, useLayoutEffect } from "react";
import { AppProvider } from "../provider";
-// Type `staticData.title` on every route so the matched-title selector
-// below stays type-safe without an inline cast.
declare module "@tanstack/react-router" {
interface StaticDataRouteOption {
title?: string;
+ titleKey?: TranslationKey;
}
}
-// Fallback while a lazy route bundle (Train/Recipes/Export) loads.
-// /chat is synchronous and never hits this.
-const RouteFallback: ReactNode = (
-
-
+
@@ -142,7 +141,7 @@ function RootLayout() {
transition={{ duration: 0.15 }}
className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"}`}
>
-
+
diff --git a/studio/frontend/src/app/routes/studio.tsx b/studio/frontend/src/app/routes/studio.tsx
index 75f1a1b937..ae7f445e94 100644
--- a/studio/frontend/src/app/routes/studio.tsx
+++ b/studio/frontend/src/app/routes/studio.tsx
@@ -15,7 +15,7 @@ const StudioPage = lazy(() =>
export const Route = createRoute({
getParentRoute: () => rootRoute,
path: "/studio",
- staticData: { title: "Train" },
+ staticData: { titleKey: "studio.routeTitle" },
beforeLoad: () => requireAuth(),
component: StudioPage,
});
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index aac5f8f8a8..4e4a6130cd 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -49,7 +49,7 @@ import {
Edit03Icon,
Globe02Icon,
HelpCircleIcon,
- Logout01Icon,
+ Logout05Icon,
Search01Icon,
PowerIcon,
PencilEdit02Icon,
@@ -90,9 +90,33 @@ import {
useTrainingRuntimeStore,
} from "@/features/training";
import type { TrainingRunSummary } from "@/features/training";
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useRef, useState, type ReactNode } from "react";
import { toast } from "@/lib/toast";
import { ShutdownDialog } from "@/components/shutdown-dialog";
+import { translate, useT, type TranslationKey } from "@/i18n";
+
+const EMPHASIS_MARKER = "__UNSLOTH_I18N_EMPHASIS_MARKER__";
+
+type AppT = ReturnType;
+
+function renderEmphasizedTranslation(
+ t: AppT,
+ key: TranslationKey,
+ emphasizedValue: string,
+): ReactNode {
+ const translated = t(key, { name: EMPHASIS_MARKER });
+ const parts = translated.split(EMPHASIS_MARKER);
+ if (parts.length === 1) return translated;
+
+ const nodes: ReactNode[] = [];
+ parts.forEach((part, index) => {
+ if (part.length > 0) nodes.push(part);
+ if (index < parts.length - 1) {
+ nodes.push({emphasizedValue});
+ }
+ });
+ return nodes;
+}
function getTourId(pathname: string): string | null {
if (pathname.startsWith("/studio")) return "studio";
@@ -185,6 +209,7 @@ function NavItem({
}
export function AppSidebar() {
+ const t = useT();
const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle();
const { pathname, search } = useRouterState({
select: (s) => ({
@@ -204,14 +229,8 @@ export function AppSidebar() {
const chatOnly = usePlatformStore((s) => s.isChatOnly());
const [shutdownOpen, setShutdownOpen] = useState(false);
- // Chat collapsible state — open by default, auto-expand on route entry
const isChatRoute = pathname.startsWith("/chat");
const isStudioRoute = pathname === "/studio" || pathname.startsWith("/studio/");
- const [chatOpen, setChatOpen] = useState(true);
- const [runsOpen, setRunsOpen] = useState(true);
-
- useEffect(() => { if (isChatRoute) setChatOpen(true); }, [isChatRoute]);
- useEffect(() => { if (isStudioRoute) setRunsOpen(true); }, [isStudioRoute]);
const scrollRef = useRef(null);
const [scrolled, setScrolled] = useState(false);
@@ -290,7 +309,7 @@ export function AppSidebar() {
try {
await renameChatItem(target.item, renameTrimmed);
} catch (err) {
- toast.error("Failed to rename chat", {
+ toast.error(translate("shell.toast.failedToRenameChat"), {
description: err instanceof Error ? err.message : undefined,
});
}
@@ -300,7 +319,7 @@ export function AppSidebar() {
const updated = await renameTrainingRun(target.run.id, nextRunDisplayName);
emitTrainingRunUpdated(updated);
} catch (err) {
- toast.error("Failed to rename run", {
+ toast.error(translate("shell.toast.failedToRenameRun"), {
description: err instanceof Error ? err.message : undefined,
});
}
@@ -320,14 +339,14 @@ export function AppSidebar() {
try {
await handleDeleteThread(target.item);
} catch (err) {
- toast.error("Failed to delete chat", {
+ toast.error(translate("shell.toast.failedToDeleteChat"), {
description: err instanceof Error ? err.message : undefined,
});
}
return;
}
if (target.run.status === "running") {
- toast.error("Cannot delete a running training run");
+ toast.error(t("shell.toast.cannotDeleteRunningRun"));
return;
}
try {
@@ -337,7 +356,7 @@ export function AppSidebar() {
}
emitTrainingRunDeleted(target.run.id);
} catch (err) {
- toast.error("Failed to delete run", {
+ toast.error(translate("shell.toast.failedToDeleteRun"), {
description: err instanceof Error ? err.message : undefined,
});
}
@@ -366,7 +385,7 @@ export function AppSidebar() {
});
}}
className="flex items-center gap-[6px] select-none"
- aria-label="Unsloth home"
+ aria-label={t("shell.aria.home")}
>
- BETA
+ {t("shell.beta")}
{!isMobile && (
@@ -387,7 +406,7 @@ export function AppSidebar() {
type="button"
onClick={togglePinned}
className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
- aria-label="Close sidebar"
+ aria-label={t("shell.aria.closeSidebar")}
>
@@ -434,7 +453,7 @@ export function AppSidebar() {
+ - Page not found + {t("shell.notFound.title")}
- {pathname} does not exist. + {t("shell.notFound.description", { path: pathname })}
- Loading...
-
-);
+function RouteFallback() {
+ const t = useT();
+
+ return (
+
+ {t("common.loading")}
+
+ );
+}
const CHAT_ONLY_ALLOWED = new Set([
"/",
@@ -68,6 +70,7 @@ const HIDDEN_NAVBAR_ROUTES = ["/onboarding", "/login", "/change-password"];
const DEFAULT_DOCUMENT_TITLE = "Unsloth Studio";
function RootLayout() {
+ const t = useT();
const pathname = useRouterState({ select: (s) => s.location.pathname });
const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname);
const isChatRoute = pathname.startsWith("/chat");
@@ -75,24 +78,20 @@ function RootLayout() {
useTrainingUnloadGuard();
- // Walk matches deepest-first; each route declares its own title.
const matchedTitle = useMatches({
select: (matches) => {
for (let i = matches.length - 1; i >= 0; i--) {
- const title = matches[i].staticData.title;
+ const { title, titleKey } = matches[i].staticData;
+ if (titleKey) return t(titleKey);
if (title) return title;
}
return null;
},
});
- // `/settings` redirects in `beforeLoad`, so its route never stays
- // matched; surface the modal's title via the store instead.
const settingsDialogOpen = useSettingsDialogStore((s) => s.open);
- const documentTitle = settingsDialogOpen ? "Settings" : matchedTitle;
+ const documentTitle = settingsDialogOpen ? t("settings.title") : matchedTitle;
- // useLayoutEffect updates the tab title before paint, avoiding a
- // one-frame flash of the previous route's title on navigation.
useLayoutEffect(() => {
document.title = documentTitle
? `${documentTitle} - ${DEFAULT_DOCUMENT_TITLE}`
@@ -116,7 +115,7 @@ function RootLayout() {
- BETA
+ {t("shell.beta")}
{!isMobile && (
@@ -387,7 +406,7 @@ export function AppSidebar() {
type="button"
onClick={togglePinned}
className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
- aria-label="Close sidebar"
+ aria-label={t("shell.aria.closeSidebar")}
>
@@ -64,11 +74,23 @@ export function ApiKeyRow({
- Created {relative(apiKey.created_at)}
+
+ {t("settings.apiKeys.created", {
+ value: relative(apiKey.created_at, t),
+ })}
+
·
- Used {relative(apiKey.last_used_at)}
+
+ {t("settings.apiKeys.used", {
+ value: relative(apiKey.last_used_at, t),
+ })}
+
·
- Expires {expiresText(apiKey.expires_at)}
+
+ {t("settings.apiKeys.expires", {
+ value: expiresText(apiKey.expires_at, t),
+ })}
+
{EXPIRY_PRESETS.map((p) => {
@@ -69,13 +74,15 @@ export function CreateKeyForm({
: "text-muted-foreground hover:text-foreground",
)}
>
- {p.label}
+ {p.value === null ? t("settings.apiKeys.never") : p.label}
);
})}
diff --git a/studio/frontend/src/features/settings/components/key-reveal-card.tsx b/studio/frontend/src/features/settings/components/key-reveal-card.tsx
index 2b589e88fe..bdcb861c1d 100644
--- a/studio/frontend/src/features/settings/components/key-reveal-card.tsx
+++ b/studio/frontend/src/features/settings/components/key-reveal-card.tsx
@@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
+import { useT } from "@/i18n";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { cn } from "@/lib/utils";
import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons";
@@ -15,6 +16,7 @@ export function KeyRevealCard({
rawKey: string;
onDone: () => void;
}) {
+ const t = useT();
const [copied, setCopied] = useState(false);
const handleCopy = async () => {
@@ -32,7 +34,7 @@ export function KeyRevealCard({
className="size-3.5 text-emerald-600 dark:text-emerald-500"
/>
- New access token created
+ {t("settings.apiKeys.newTokenCreated")}
- Copy now — this won't be shown again. + {t("settings.apiKeys.copyNow")}