diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 673b2f3cc5..2edcae8ab2 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -285,7 +285,15 @@ jobs: # The PR-time CI must validate the code in this PR; PyPI unsloth # may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py # (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream. - pip install --no-deps unsloth_zoo + # unsloth_zoo from git main mirrors every other CI (Core / MLX / + # install.sh) so PR-time validation sees the same zoo HEAD. + for attempt in 1 2 3; do + if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install --no-deps -e ./unsloth - name: Convert notebooks for AST scan diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 63eb70f7f1..ee5bbe8633 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -144,9 +144,19 @@ jobs: # versions ship a CPU build that imports cleanly on Linux. pip install 'bitsandbytes>=0.45' # unsloth.device_type imports unsloth_zoo.utils.Version at module - # scope, so the conftest preload needs unsloth_zoo even though - # it is an optional dep of unsloth. - pip install 'unsloth_zoo>=2026.5.1' + # scope, so the conftest preload needs unsloth_zoo. Pull from + # git main so this job sees the same zoo HEAD as Core / MLX / + # install.sh do (otherwise a fix on zoo main hides until release). + # No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'` + # behaviour so triton etc. still come in for the Repo tests CPU + # collection imports. + for attempt in 1 2 3; do + if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then + break + fi + [ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; } + sleep $((5 * attempt)) + done pip install -e . --no-deps - name: Repo tests (CPU, auto-discovered) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9d80fe6ff5..1919fac9c5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.13 + rev: v0.15.14 hooks: - id: ruff args: diff --git a/README.md b/README.md index 3699c50736..ecf0f8a7a3 100644 --- a/README.md +++ b/README.md @@ -215,6 +215,9 @@ Then to launch every time: unsloth studio -p 8888 ``` +#### Advanced launch options +Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. Explicit `OMP_NUM_THREADS` / `MKL_NUM_THREADS` / `OPENBLAS_NUM_THREADS` / `NUMEXPR_NUM_THREADS` still take precedence. + #### Uninstall The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows): diff --git a/install.ps1 b/install.ps1 index b26566cc3d..c4e8d8d522 100644 --- a/install.ps1 +++ b/install.ps1 @@ -887,12 +887,19 @@ shell.Run cmd, 0, False } # ── Check winget ── + # winget is only needed to install Python or uv. If both are + # already on PATH (Windows ARM64 GitHub-hosted runners, manual + # python.org + Astral uv installs, corporate locked-down hosts + # without the Store, etc.) the script can proceed without it. + # We defer the hard failure to the Python / uv install branches + # below, where winget is actually invoked. Write-TauriLog "STEP" "Checking system dependencies" - if (-not (Get-Command winget -ErrorAction SilentlyContinue)) { - step "winget" "not available" "Red" - substep "Install it from https://aka.ms/getwinget" "Yellow" - substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow" - return (Exit-InstallFailure "winget is not available") + $script:WingetAvailable = [bool](Get-Command winget -ErrorAction SilentlyContinue) + if ($script:WingetAvailable) { + step "winget" "available" + } else { + step "winget" "not available -- will require Python + uv to be already installed" "Yellow" + substep "Get it from https://aka.ms/getwinget if Python / uv are not already on PATH." "Yellow" } # ── Helper: detect a working Python 3.11-3.13 on the system ── @@ -973,6 +980,12 @@ shell.Run cmd, 0, False step "python" "Python $($DetectedPython.Version) already installed" } if (-not $DetectedPython) { + if (-not $script:WingetAvailable) { + Write-Host "[ERROR] No compatible Python (3.11-3.13) found and winget is unavailable on this host." -ForegroundColor Red + Write-Host " Install Python $PythonVersion from https://www.python.org/downloads/" -ForegroundColor Yellow + Write-Host " and re-run this installer (make sure 'Add Python to PATH' is checked)." -ForegroundColor Yellow + return (Exit-InstallFailure "winget required to install Python on this host") + } substep "installing Python ${PythonVersion}..." $pythonPackageId = "Python.Python.$PythonVersion" # Temporarily lower ErrorActionPreference so that winget stderr @@ -1024,14 +1037,19 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing uv package manager" if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { substep "installing uv package manager..." - $prevEAP = $ErrorActionPreference - $ErrorActionPreference = "Continue" - try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} - $ErrorActionPreference = $prevEAP - Refresh-SessionPath - # Fallback: if winget didn't put uv on PATH, try the PowerShell installer + if ($script:WingetAvailable) { + $prevEAP = $ErrorActionPreference + $ErrorActionPreference = "Continue" + try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {} + $ErrorActionPreference = $prevEAP + Refresh-SessionPath + } + # Fallback: if winget is unavailable or didn't put uv on PATH, + # use Astral's official PowerShell installer. This is the only + # supported path on hosts without winget (Windows ARM64 runners, + # corporate machines without the Store, etc.). if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { - substep "trying alternative uv installer..." "Yellow" + substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow" Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1") Refresh-SessionPath } @@ -1235,7 +1253,10 @@ shell.Run cmd, 0, False if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = & $NvidiaSmiExe 2>&1 | Out-String - if ($output -match 'CUDA Version:\s+(\d+)\.(\d+)') { + # Newer NVIDIA drivers (e.g. 610.x on Windows) print + # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". + # Accept both spellings so we don't fall through to the cu126 default. + if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') { $major = [int]$Matches[1]; $minor = [int]$Matches[2] if ($major -ge 13) { return "$baseUrl/cu130" } if ($major -eq 12 -and $minor -ge 8) { return "$baseUrl/cu128" } @@ -1300,15 +1321,11 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } if ($baseInstallExit -eq 0) { - # Install pydantic WITH deps so pip pins pydantic-core to - # the exact version pydantic's metadata requires. The - # --no-deps install of no-torch-runtime.txt below would - # otherwise pick the latest of each independently and - # trip pydantic's _ensure_pydantic_core_version check. - # pydantic's deps (annotated-types, pydantic-core, - # typing-extensions, typing-inspection) are torch-free. + # Resolve pydantic WITH deps so pip pins pydantic-core + # to the matching version (no-torch-runtime.txt below + # is --no-deps). All transitive deps are torch-free. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { @@ -1318,7 +1335,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1356,10 +1373,9 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } if ($baseInstallExit -eq 0) { - # Install pydantic WITH deps so pip pins pydantic-core to - # the matching version (see migrated branch above). + # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { @@ -1369,7 +1385,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.6" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1397,7 +1413,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.6" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -1631,6 +1647,33 @@ shell.Run cmd, 0, False # New-StudioShortcuts gates the .lnk shortcuts on env-mode internally. New-StudioShortcuts -UnslothExePath $UnslothExe + # Warn if another 'unsloth' wins on PATH (different venv, system pip). + # Mirrors install.sh; absolute path is still the most reliable launch. + # Uses content-hash equality (Get-FileHash) so hardlinks, symlinks, and + # identical copies of the installer's shim don't false-trigger. CommandType + # Application restricts the probe to real executables (skips aliases, + # functions, scripts). + try { + $_pathCmd = Get-Command unsloth -CommandType Application -ErrorAction SilentlyContinue | Select-Object -First 1 + if ($_pathCmd) { + $_pathExe = $_pathCmd.Source + $_installedHash = (Get-FileHash -LiteralPath $UnslothExe -Algorithm SHA256 -ErrorAction SilentlyContinue).Hash + $_pathHash = (Get-FileHash -LiteralPath $_pathExe -Algorithm SHA256 -ErrorAction SilentlyContinue).Hash + if ($_installedHash -and $_pathHash -and ($_installedHash -ne $_pathHash)) { + Write-Host "" + step "warning" "another 'unsloth' wins on PATH:" "Yellow" + substep $_pathExe + substep "this installer's binary is at:" + substep $UnslothExe + substep "to use this install, call the absolute path above," + substep "or put its dir earlier on PATH." + Write-Host "" + } + } + } catch { + # Diagnostic only; never block install on a probe failure. + } + # In interactive terminals, ask the user before starting Studio. # In non-interactive environments (CI, Docker) just print instructions. $IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) diff --git a/install.sh b/install.sh index 9bdd935171..a3b76a3011 100755 --- a/install.sh +++ b/install.sh @@ -1290,6 +1290,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi +# Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file). +if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" + if [ -f "$_OVERRIDES_FILE" ]; then + export UV_OVERRIDE="$_OVERRIDES_FILE" + fi +fi + _TAURI_INITIAL_GPU_BRANCH="unknown" if [ "$SKIP_TORCH" = true ]; then _TAURI_INITIAL_GPU_BRANCH="no_torch" @@ -1675,9 +1683,15 @@ get_torch_index_url() { fi echo "$_base/cpu"; return fi - # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P) + # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). + # Newer NVIDIA drivers (e.g. 610.x) print "CUDA UMD Version: X.Y" instead + # of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions + # (POSIX sed does not support "?" without -E). The two patterns are + # mutually exclusive per line, so head -1 picks the first emitted match. _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ - | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ + | sed -n \ + -e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ + -e 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ | head -1) if [ -z "$_cuda_ver" ]; then echo "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" >&2 @@ -1865,14 +1879,10 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.6" unsloth-zoo - # Install pydantic WITH deps so pip pins pydantic-core to the - # exact version pydantic's own metadata requires. The --no-deps - # install below would otherwise pick the latest of each - # independently and trip pydantic's _ensure_pydantic_core_version - # check on the next import. pydantic's deps (annotated-types, - # pydantic-core, typing-extensions, typing-inspection) are - # torch-free, so this is safe on the no-torch path. + "unsloth>=2026.5.8" unsloth-zoo + # Resolve pydantic WITH deps so pip pins pydantic-core to the + # matching version (no-torch-runtime.txt below is --no-deps). + # All transitive deps are torch-free. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" @@ -1882,7 +1892,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.6" unsloth-zoo + "unsloth>=2026.5.8" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2050,9 +2060,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.5.6" unsloth-zoo - # Install pydantic WITH deps so pip pins pydantic-core to the - # exact version pydantic requires (see migrated branch above). + "unsloth>=2026.5.8" unsloth-zoo + # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" @@ -2069,7 +2078,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.5.6" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2101,7 +2110,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.6" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2113,12 +2122,6 @@ else fi fi -# ── Install mlx-vlm on Apple Silicon (optional, for VLM training) ── -if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then - substep "installing mlx-vlm (VLM training support)..." - run_install_cmd "install mlx-vlm" uv pip install --python "$_VENV_PY" mlx-vlm -fi - # ── Run studio setup ── tauri_log "STEP" "Running Studio setup" # When --local, use the repo's own setup.sh directly. @@ -2268,6 +2271,38 @@ if [ "$TAURI_MODE" = true ]; then exit 0 fi +# Warn if another 'unsloth' wins on PATH (different venv, system pip, etc). +# Users typing `unsloth studio` later would hit that binary instead of the +# one just installed; the runtime now falls back via UNSLOTH_STUDIO_HOME +# but the absolute path is still the most reliable launch. +# Uses the venv python (just created above) for path canonicalization so +# this works on macOS (BSD readlink has no -f) as well as Linux/WSL. +_installed_bin="$VENV_DIR/bin/unsloth" +_path_unsloth=$(command -v unsloth 2>/dev/null || true) +if [ -n "$_path_unsloth" ] && [ -x "$VENV_DIR/bin/python" ]; then + # Canonicalize via the venv python (BSD readlink lacks -f on macOS). + # If either side fails to resolve, skip the check entirely rather than + # comparing raw paths (which would false-trigger on symlink targets). + _canon() { + "$VENV_DIR/bin/python" -c \ + 'import os, sys; print(os.path.realpath(sys.argv[1]))' \ + "$1" 2>/dev/null + } + _installed_real=$(_canon "$_installed_bin") + _path_real=$(_canon "$_path_unsloth") + if [ -n "$_installed_real" ] && [ -n "$_path_real" ] \ + && [ "$_installed_real" != "$_path_real" ]; then + echo "" + step "warning" "another 'unsloth' wins on PATH:" "$C_WARN" + substep "$_path_unsloth" + substep "this installer's binary is at:" + substep "$_installed_bin" + substep "to use this install, run the absolute path above," + substep "alias unsloth, or put its dir earlier on PATH." + echo "" + fi +fi + echo "" printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" printf " ${C_DIM}%s${C_RST}\n" "$RULE" diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 7fff36aefd..b4ec0ccd94 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -109,7 +109,7 @@ def _apply_data_designer_image_context_patch() -> None: return try: - from data_designer.config.models import ImageContext + from data_designer.config.models import ImageContext # pyright: ignore[reportMissingImports] except ImportError: return @@ -131,7 +131,7 @@ def _apply_data_designer_image_context_patch() -> None: def build_model_providers(recipe: dict[str, Any]): - from data_designer.config.models import ModelProvider + from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] providers: list[ModelProvider] = [] for provider in recipe.get("model_providers", []): @@ -174,7 +174,7 @@ def _validate_recipe_runtime_support( def build_mcp_providers( recipe: dict[str, Any], ) -> list: - from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider + from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] providers: list[MCPProvider | LocalStdioMCPProvider] = [] for provider in recipe.get("mcp_providers", []): @@ -214,16 +214,42 @@ def build_mcp_providers( return providers +def _strip_frontend_model_config_metadata(recipe: dict[str, Any]) -> dict[str, Any]: + model_configs = recipe.get("model_configs") + if not isinstance(model_configs, list): + return recipe + + changed = False + next_model_configs: list[Any] = [] + for model_config in model_configs: + if isinstance(model_config, dict) and "gguf_variant" in model_config: + next_model_config = dict(model_config) + next_model_config.pop("gguf_variant", None) + next_model_configs.append(next_model_config) + changed = True + continue + next_model_configs.append(model_config) + + if not changed: + return recipe + + return { + **recipe, + "model_configs": next_model_configs, + } + + def build_config_builder(recipe: dict[str, Any]): _apply_data_designer_image_context_patch() - from data_designer.config import DataDesignerConfigBuilder - from data_designer.config.processors import ProcessorType + from data_designer.config import DataDesignerConfigBuilder # pyright: ignore[reportMissingImports] + from data_designer.config.processors import ProcessorType # pyright: ignore[reportMissingImports] recipe_core = { key: value for key, value in recipe.items() if key not in {"model_providers", "mcp_providers"} } + recipe_core = _strip_frontend_model_config_metadata(recipe_core) recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators( recipe_core ) @@ -256,8 +282,9 @@ def create_data_designer( artifact_path: str | None = None, ): _apply_data_designer_image_context_patch() - from data_designer.interface.data_designer import DataDesigner + from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports] + recipe = _strip_frontend_model_config_metadata(recipe) model_providers = build_model_providers(recipe) _validate_recipe_runtime_support(recipe, model_providers) @@ -265,7 +292,7 @@ def create_data_designer( # when the pipeline contains no LLM columns. Supply a lightweight stub # so sampler/expression-only recipes can run without a real provider. if not model_providers: - from data_designer.config.models import ModelProvider + from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] model_providers = [ ModelProvider( diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 4ab95d896f..7cabd382eb 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -475,6 +475,7 @@ class ExportBackend: self.current_model.save_pretrained_merged( save_directory, self.current_tokenizer, + save_method = "merged_16bit", ) else: self.current_model.save_pretrained(save_directory) @@ -510,6 +511,7 @@ class ExportBackend: self.current_model.save_pretrained_merged( tmp_dir, self.current_tokenizer, + save_method = "merged_16bit", ) self.current_model.push_to_hub_merged( repo_id, diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 25e1725337..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 @@ -67,6 +69,180 @@ _ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile( r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)" ) _OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)") +_OPENAI_REASONING_STATUSES = {"in_progress", "completed", "incomplete"} + + +def _openai_image_replay_requires_reasoning(model: str) -> bool: + normalized = model.strip().lower() + return normalized.startswith("gpt-5") or normalized.startswith("o") + + +def _sanitize_openai_reasoning_replay_item( + item: Any, +) -> Optional[dict[str, Any]]: + """Return a Responses input-safe reasoning item, if ``item`` is one. + + OpenAI's image-generation docs allow follow-up edits by sending the + previous ``image_generation_call`` id. Reasoning models can additionally + require the paired ``reasoning`` output item in manually managed context, + so keep the public replay fields only and drop everything else. + """ + if not isinstance(item, dict) or item.get("type") != "reasoning": + return None + item_id = item.get("id") + if not isinstance(item_id, str) or not item_id: + return None + summary_parts: list[dict[str, str]] = [] + summary = item.get("summary") + if isinstance(summary, list): + for part in summary: + if not isinstance(part, dict): + continue + if part.get("type") != "summary_text": + continue + text = part.get("text") + if isinstance(text, str): + summary_parts.append({"type": "summary_text", "text": text}) + replay_item: dict[str, Any] = { + "type": "reasoning", + "id": item_id, + "summary": summary_parts, + } + status = item.get("status") + if isinstance(status, str) and status in _OPENAI_REASONING_STATUSES: + replay_item["status"] = status + return replay_item + + +# OpenAI Responses inline citation markers: `citeSOURCE_ID[id2...][LOCATOR]` +# using private-use codepoints (see +# https://developers.openai.com/api/docs/guides/citation-formatting). +# Group 1 holds the delim-separated tokens; each resolvable token expands +# to `[[N]](URL)`, unresolved tokens (locators, unknown ids) drop silently +# so no garbled glyph reaches the renderer. +_OPENAI_CITE_OPEN = "cite" +_OPENAI_CITE_STOP = "" +_OPENAI_CITE_DELIM = "" +_OPENAI_CITATION_MARKER = re.compile( + f"{_OPENAI_CITE_OPEN}([^{_OPENAI_CITE_STOP}]+){_OPENAI_CITE_STOP}" +) + + +def _build_citation_lookup( + url_citations: list[dict[str, Any]], +) -> dict[str, tuple[int, str]]: + """Map every known ``source_id`` alias to ``(citation_index, url)``. + + Accepts singular ``source_id`` and plural ``source_ids``. First-seen + wins on alias collision so an earlier citation keeps its number. + """ + by_source: dict[str, tuple[int, str]] = {} + for idx, cit in enumerate(url_citations, start = 1): + url = cit.get("url") + if not isinstance(url, str) or not url: + continue + aliases: list[str] = [] + sid = cit.get("source_id") + if isinstance(sid, str) and sid: + aliases.append(sid) + sids = cit.get("source_ids") + if isinstance(sids, list): + aliases.extend(s for s in sids if isinstance(s, str) and s) + for alias in aliases: + by_source.setdefault(alias, (idx, url)) + return by_source + + +def _replace_openai_citation_markers( + text: str, + url_citations: list[dict[str, Any]], +) -> str: + """Rewrite `\\ue200cite\\ue202SOURCE_ID[\\ue202LOCATOR]\\ue201` markers into + `[[N]](URL)` per resolvable id. Multi-source markers expand to one link + per id; unresolved tokens drop silently. Idempotent on text without + private-use codepoints. + """ + if not text or _OPENAI_CITE_STOP not in text: + return text + by_source = _build_citation_lookup(url_citations) + + def _sub(match: re.Match[str]) -> str: + # Try every delim-split token; unresolved tokens drop silently. + # Handles multi-source (all resolve) and source+locator (only the + # id resolves, locator drops). Empty result strips the marker. + rendered: list[str] = [] + for tok in match.group(1).split(_OPENAI_CITE_DELIM): + if not tok: + continue + hit = by_source.get(tok) + if hit is None: + continue + idx, url = hit + rendered.append(f"[[{idx}]]({url})") + return "".join(rendered) + + return _OPENAI_CITATION_MARKER.sub(_sub, text) + + +def _rewrite_citation_markers_partial( + text: str, + url_citations: list[dict[str, Any]], +) -> tuple[str, bool]: + """Like ``_replace_openai_citation_markers`` but also reports whether + any marker referenced a source_id not yet in ``url_citations``. + + The ``annotation.added`` event for a url_citation typically arrives + AFTER the delta carrying the marker referencing it. Callers buffer the + segment until a later event records the annotation; unresolved markers + are left verbatim so a follow-up pass still parses cleanly. + """ + if not text or _OPENAI_CITE_STOP not in text: + return text, False + by_source = _build_citation_lookup(url_citations) + has_unresolved = False + + def _sub(match: re.Match[str]) -> str: + nonlocal has_unresolved + tokens = [t for t in match.group(1).split(_OPENAI_CITE_DELIM) if t] + rendered: list[str] = [] + any_unresolved = False + for tok in tokens: + hit = by_source.get(tok) + if hit is None: + any_unresolved = True + continue + idx, url = hit + rendered.append(f"[[{idx}]]({url})") + # Leave the whole marker verbatim if any token is unresolved so the + # caller can re-run once the late annotation lands; partial emission + # would lose the unresolved ids once the source text is dropped. + if any_unresolved: + has_unresolved = True + return match.group(0) + return "".join(rendered) + + return _OPENAI_CITATION_MARKER.sub(_sub, text), has_unresolved + + +def _split_pending_citation_tail(text: str) -> tuple[str, str]: + """Split ``text`` into ``(head, pending_tail)`` for streamed deltas. + + A citation marker can straddle two SSE deltas (e.g. delta-1 ends with + ``\\ue200citetu`` and delta-2 starts with ``rn0view0\\ue201``); the + unterminated tail is buffered and prepended onto the next delta so the + rewriter sees a complete marker. ``pending_tail`` is the longest suffix + starting with ``\\ue200`` and lacking ``\\ue201``; ``head`` is safe to + emit. Empty tail when ``text`` has no open marker or a fully closed one. + """ + if not text: + return text, "" + last_open = text.rfind("") + if last_open == -1: + return text, "" + # Stop byte after the last open byte means the marker closed in this delta. + if _OPENAI_CITE_STOP in text[last_open:]: + return text, "" + return text[:last_open], text[last_open:] class _AnthropicThinkingSpec(NamedTuple): @@ -173,10 +349,87 @@ _ANTHROPIC_COMPACTION_TYPE = "compact_20260112" _ANTHROPIC_COMPACTION_MIN = 50_000 +# Anthropic fast-mode beta (Opus 4.6 / 4.7 only, per +# https://platform.claude.com/docs/en/build-with-claude/fast-mode). +# Mutually exclusive with the Priority service tier. +_ANTHROPIC_FAST_MODE_BETA = "fast-mode-2026-02-01" +_ANTHROPIC_FAST_MODE_PREFIXES = ( + "claude-opus-4-7", + "claude-opus-4-6", +) + + def _anthropic_supports_compaction(model: str) -> bool: return model.startswith(_ANTHROPIC_COMPACTION_PREFIXES) +def _anthropic_supports_fast_mode(model: str) -> bool: + # Require a family boundary ("" or "-") after the prefix so IDs like + # "claude-opus-4-70" / "claude-opus-4-7b" do not match. + return any( + model == p or model.startswith(f"{p}-") for p in _ANTHROPIC_FAST_MODE_PREFIXES + ) + + +# Cap on ``cited_text`` forwarded in document_citations tool_events; +# keeps SSE bytes bounded on multi-KB cited spans (frontend trims to +# 240 chars anyway). +_CITED_TEXT_MAX_LEN = 512 + + +def _anthropic_citation_key(citation: dict[str, Any]) -> tuple: + """Stable dedup key for an Anthropic ``citations_delta.citation``. + + Anchor fields vary per type (char_location, page_location, + content_block_location, search_result_location); both start AND + exclusive end indices are part of the key so same-start / + different-end pairs stay distinct. search_result_location keys on + ``search_result_index`` + ``source`` instead of document_index so + distinct results with the same source don't collapse. Unknown + shapes fall back to a stringified copy (more entries, never + collisions). See + https://platform.claude.com/docs/en/build-with-claude/citations + and https://platform.claude.com/docs/en/build-with-claude/search-results. + """ + ctype = citation.get("type") + doc = citation.get("document_index") + title = citation.get("document_title") or "" + if ctype == "char_location": + return ( + ctype, + doc, + title, + citation.get("start_char_index"), + citation.get("end_char_index"), + ) + if ctype == "page_location": + return ( + ctype, + doc, + title, + citation.get("start_page_number"), + citation.get("end_page_number"), + ) + if ctype == "content_block_location": + return ( + ctype, + doc, + title, + citation.get("start_block_index"), + citation.get("end_block_index"), + ) + if ctype == "search_result_location": + return ( + ctype, + citation.get("search_result_index"), + citation.get("source"), + citation.get("title") or "", + citation.get("start_block_index"), + citation.get("end_block_index"), + ) + return (ctype, _json.dumps(citation, sort_keys = True)) + + class _MistralThinkingSpec(NamedTuple): models: tuple[str, ...] style: Literal["prompt_mode", "reasoning_effort", "disabled"] @@ -255,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, @@ -295,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]: @@ -315,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". @@ -329,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( @@ -343,11 +861,14 @@ 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]: """ @@ -360,8 +881,39 @@ class ExternalProviderClient: supplies a value the provider accepts — the frontend's provider-capability map already filters these per provider, so we treat them as opt-in here. + + ``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, @@ -376,6 +928,8 @@ class ExternalProviderClient: anthropic_code_exec_container_id, prompt_cache_ttl, compaction_threshold, + tool_choice, + fast_mode = fast_mode, ): yield line return @@ -398,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 ): @@ -438,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. @@ -480,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"): @@ -503,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 @@ -521,11 +1067,19 @@ class ExternalProviderClient: plugins.append({"id": "web"}) body["plugins"] = plugins logger.info( - "OpenRouter web_search: attached plugins=[{id: 'web'}] " - "(model=%s)", + "OpenRouter web_search: attached plugins=[{id: 'web'}] (model=%s)", 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)", @@ -561,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 []) ) @@ -594,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", @@ -849,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", @@ -1186,6 +1734,9 @@ 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]: """ Call the Anthropic Messages API and translate its SSE to OpenAI format. @@ -1214,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:...}` @@ -1305,6 +1892,11 @@ class ExternalProviderClient: "media_type": media_type, "data": b64data, }, + # Opt into Anthropic's natural-citation + # pipeline; without this no citations_delta + # events fire. See + # https://platform.claude.com/docs/en/build-with-claude/citations + "citations": {"enabled": True}, } if title: doc_block["title"] = title @@ -1316,10 +1908,42 @@ class ExternalProviderClient: "type": "url", "url": url, }, + "citations": {"enabled": True}, } 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 @@ -1327,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. @@ -1353,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", @@ -1486,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( { @@ -1509,25 +2191,18 @@ class ExternalProviderClient: ) body["tools"] = anthropic_tools - # Anthropic server-side web_fetch — see - # https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-fetch-tool - # `web_fetch_20250910` reads a single URL (text or PDF) and - # returns a document block in a `web_fetch_tool_result`. For - # safety Anthropic only lets the model fetch URLs that already - # appeared in the conversation (user message, prior tool - # result, web_search hit) — there is no domain restriction we - # have to apply locally. No beta header is required today; the - # tool ships under the standard `2023-06-01` API version. We - # mirror the web_search wiring: max_uses cap, opt in via - # `enabled_tools=["web_fetch"]`, citations off by default - # because the frontend already paints source pills from the - # generic tool_end payload. - 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( { - "type": "web_fetch_20250910", + "type": _anthropic_web_fetch_version(model), "name": "web_fetch", "max_uses": 5, } @@ -1535,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 []) @@ -1563,38 +2227,20 @@ 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 and _anthropic_supports_compaction(model) ) - if compaction_active: + if compaction_active and compaction_threshold is not None: trigger_value = max( int(compaction_threshold), _ANTHROPIC_COMPACTION_MIN, @@ -1611,13 +2257,18 @@ class ExternalProviderClient: ] } + # fast_mode is Opus 4.6/4.7 only; silently drop elsewhere. + # Incompatible with the Priority service_tier (frontend gate + # prevents both at once; backend lets Anthropic 400 if combined). + fast_mode_active = bool(fast_mode) and _anthropic_supports_fast_mode(model) + if fast_mode_active: + body["speed"] = "fast" + 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, " @@ -1632,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", @@ -1654,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()] @@ -1669,6 +2310,8 @@ class ExternalProviderClient: beta_parts.append(_ANTHROPIC_CODE_EXECUTION_BETA) if compaction_active and _ANTHROPIC_COMPACTION_BETA not in beta_parts: beta_parts.append(_ANTHROPIC_COMPACTION_BETA) + if fast_mode_active and _ANTHROPIC_FAST_MODE_BETA not in beta_parts: + beta_parts.append(_ANTHROPIC_FAST_MODE_BETA) if beta_parts: request_headers["anthropic-beta"] = ",".join(beta_parts) @@ -1722,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]] = {} @@ -1767,6 +2397,12 @@ class ExternalProviderClient: # the next turn. current_compaction: Optional[dict[str, Any]] = None compaction_blocks_seen = 0 + # Document citations from ``citations_delta`` events. + # Deduped by type-specific anchor key; inline [N] is + # injected after each cited run, and the full list is + # forwarded as a synthetic document_citations tool_event + # on message_stop for the Sources panel. + document_citations: list[dict[str, Any]] = [] # Counts surfaced in the final log line so reports of # "Code execution did nothing" can be triaged at a # glance. generated_files_count is interesting for the @@ -1806,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", @@ -2063,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 "", @@ -2084,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: @@ -2118,10 +2740,27 @@ class ExternalProviderClient: thinking_open = False if text: yield _content_chunk(text) - # Citations on text deltas are attached - # per-call by Anthropic via the - # `web_search_tool_result` block; we don't - # need to scrape them off the text events. + # web_search citations: web_search_tool_result. + # User-doc citations: citations_delta below. + elif delta_type == "citations_delta": + # One citation per event; collapse onto a + # numbered footnote list and inject [N] + # inline. See + # https://platform.claude.com/docs/en/build-with-claude/citations + cit = delta.get("citation") + if isinstance(cit, dict): + key = _anthropic_citation_key(cit) + idx_for_marker: Optional[int] = None + for idx, existing in enumerate( + document_citations, start = 1 + ): + if existing.get("_key") == key: + idx_for_marker = idx + break + if idx_for_marker is None: + document_citations.append({**cit, "_key": key}) + idx_for_marker = len(document_citations) + yield _content_chunk(f"[{idx_for_marker}]") elif delta_type == "input_json_delta": # Streamed partial_json carrying tool inputs # — the search query for web_search, or the @@ -2342,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 @@ -2410,6 +3039,29 @@ class ExternalProviderClient: # finish_reason="stop" chunk that would # truncate the rendered message in the UI. mapped = _finish_reason_map.get(stop_reason, "stop") + # Streaming refusal: emit a visible notice + # plus an out-of-band _toolEvent so the + # frontend can prune the refused turn. + # The mapped finish_reason is + # "content_filter" per OpenAI spec. + # https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals + if stop_reason == "refusal": + logger.warning( + "Anthropic refusal stop_reason (model=%s)", + model, + ) + # Drop signal rides _toolEvent (not + # text) so assistant content cannot + # spoof a context reset. + yield _content_chunk( + "\n\n_The response was stopped by " + "Anthropic's safety classifier. Edit " + "or remove the previous turn and try " + "again._" + ) + yield _emit_tool_event( + {"type": "anthropic_refusal"} + ) if mapped is not None: chunk = { "id": completion_id, @@ -2428,6 +3080,29 @@ class ExternalProviderClient: if thinking_open: yield _content_chunk("") thinking_open = False + # Forward document_citations so the Sources + # panel can render the inline [N] footnotes. + # ``cited_text`` is truncated server-side to + # keep SSE bytes bounded on long spans. + if document_citations: + clean_cits = [] + for c in document_citations: + entry = {k: v for k, v in c.items() if k != "_key"} + cited = entry.get("cited_text") + if ( + isinstance(cited, str) + and len(cited) > _CITED_TEXT_MAX_LEN + ): + entry["cited_text"] = ( + cited[:_CITED_TEXT_MAX_LEN] + "…" + ) + clean_cits.append(entry) + yield _emit_tool_event( + { + "type": "document_citations", + "citations": clean_cits, + } + ) # Final include_usage-style chunk so callers can # see cache_creation / cache_read without # scraping the server log. @@ -2545,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]], @@ -2558,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 @@ -2572,10 +4965,23 @@ class ExternalProviderClient: """ import json as _json + is_openai_cloud = _is_openai_family_cloud(self.base_url) + image_generation_requested = bool( + enabled_tools and "image_generation" in enabled_tools and is_openai_cloud + ) + # Split system messages out into a single `instructions` string and # 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: role = msg.get("role") content = msg.get("content", "") @@ -2590,12 +4996,133 @@ 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 if isinstance(content, list): translated_parts: list[dict[str, Any]] = [] + used_previous_response_id = False for part in content: part_type = part.get("type") if part_type == "text": @@ -2610,6 +5137,36 @@ class ExternalProviderClient: translated_parts.append( {"type": "input_image", "image_url": url} ) + elif ( + part_type == "reasoning" + and role == "assistant" + and image_generation_requested + ): + replay_item = _sanitize_openai_reasoning_replay_item(part) + if replay_item: + openai_replay_items.append(replay_item) + elif ( + part_type == "image_generation_call" + and role == "assistant" + and image_generation_requested + ): + response_id = ( + part.get("response_id") + or part.get("openai_response_id") + or part.get("previous_response_id") + ) + call_id = part.get("id") or part.get("image_generation_call_id") + if isinstance(call_id, str) and call_id: + if isinstance(response_id, str) and response_id: + previous_response_id = response_id + input_items = [] + translated_parts = [] + used_previous_response_id = True + else: + previous_response_id = None + openai_replay_items.append( + {"type": "image_generation_call", "id": call_id} + ) elif part_type == "input_document": # OpenAI Responses accepts PDFs / docs as # `{type:"input_file", file_data:"data:application/pdf;base64,..."}` @@ -2647,9 +5204,59 @@ class ExternalProviderClient: if filename: block["filename"] = filename translated_parts.append(block) - if translated_parts: + if translated_parts and not used_previous_response_id: input_items.append({"role": role, "content": translated_parts}) + if previous_response_id: + # OpenAI's documented multi-turn image generation path can use + # `previous_response_id` to carry the prior generated image and + # paired reasoning state. Prefer that over manual item replay when + # we captured the response id; keep replay below as a fallback for + # older stored turns that only have an image_generation_call id. + openai_replay_items = [] + elif ( + _openai_image_replay_requires_reasoning(model) + and reasoning_effort != "none" + and enable_thinking is not False + ): + filtered_replay_items: list[dict[str, Any]] = [] + has_reasoning_replay = False + dropped_image_replay_without_reasoning = False + for item in openai_replay_items: + if item.get("type") == "reasoning": + has_reasoning_replay = True + filtered_replay_items.append(item) + elif item.get("type") == "image_generation_call": + if has_reasoning_replay: + filtered_replay_items.append(item) + else: + dropped_image_replay_without_reasoning = True + else: + filtered_replay_items.append(item) + openai_replay_items = filtered_replay_items + if dropped_image_replay_without_reasoning: + yield _error_sse_line( + 400, + "OpenAI image edit reference is missing paired reasoning state. " + "Regenerate the image, then retry the edit.", + self.provider_type, + ) + return + image_generation_has_reference = bool( + previous_response_id + or any( + isinstance(item, dict) and item.get("type") == "image_generation_call" + for item in openai_replay_items + ) + ) + if openai_replay_items: + insert_at = len(input_items) + for index in range(len(input_items) - 1, -1, -1): + if input_items[index].get("role") == "user": + insert_at = index + break + input_items[insert_at:insert_at] = openai_replay_items + # NOTE: gpt-5.x / o3 / gpt-4.5 are reasoning-class models. They reject # temperature and top_p with `Unsupported parameter` 400s on # /v1/responses (and on /v1/chat/completions for the same families). @@ -2665,6 +5272,8 @@ class ExternalProviderClient: "input": input_items, "stream": True, } + if previous_response_id: + body["previous_response_id"] = previous_response_id # `summary: "auto"` is what makes /v1/responses emit reasoning # summary events — without it OpenAI returns no thinking text on # most reasoning models, the SSE handler has no … @@ -2696,45 +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. - is_openai_cloud = _is_openai_family_cloud(self.base_url) + # 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 @@ -2747,46 +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 = bool( enabled_tools and "image_generation" in enabled_tools and is_openai_cloud ) - if enabled_tools: - tools_array: list[dict[str, Any]] = [] - if "web_search" in enabled_tools: + + def _openai_image_generation_tool() -> dict[str, Any]: + tool: dict[str, Any] = {"type": "image_generation"} + if image_generation_has_reference: + # Force edit mode so the prior call id is used as context. + tool["action"] = "edit" + return tool + + # 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 = { @@ -2796,10 +5421,12 @@ class ExternalProviderClient: else: shell_env = {"type": "container_auto"} tools_array.append({"type": "shell", "environment": shell_env}) - if image_generation_enabled_openai: - tools_array.append({"type": "image_generation"}) + 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('/', '-')}" @@ -2813,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", @@ -2828,12 +5463,17 @@ class ExternalProviderClient: tools_array_attempt.append( {"type": "shell", "environment": env_attempt} ) - if image_generation_enabled_openai: - tools_array_attempt.append({"type": "image_generation"}) + 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: @@ -2895,56 +5535,106 @@ 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, str]] = [] - # 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} } + all_url_citations: list[dict[str, Any]] = [] + # 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 + last_openai_reasoning_replay_item: Optional[dict[str, Any]] = None + openai_reasoning_replay_items: dict[str, dict[str, Any]] = {} + image_generation_calls_started: set[str] = set() + # Buffer for a citation marker straddling two delta events; + # prepended onto the next delta. See _split_pending_citation_tail. + pending_marker_tail: str = "" + # Segments deferred while their markers reference unseen + # source_ids; held in arrival order so output never + # leapfrogs an earlier deferred segment. Flushed on + # annotation events and force-flushed at end-of-stream + # with leftover private-use codepoints stripped. + pending_citation_segments: list[str] = [] + + def _record_openai_response_id(payload: dict[str, Any]) -> None: + nonlocal current_openai_response_id + response_obj = payload.get("response") + candidates: list[Any] = [] + if isinstance(response_obj, dict): + candidates.append(response_obj.get("id")) + candidates.append(payload.get("response_id")) + for candidate in candidates: + if isinstance(candidate, str) and candidate: + current_openai_response_id = candidate + return + + def _drain_pending_segments(force: bool) -> str: + """Re-attempt resolution on buffered segments in order. + Stops at the first still-unresolved segment unless + ``force`` (end-of-stream), where lingering markers are stripped.""" + out: list[str] = [] + while pending_citation_segments: + seg = pending_citation_segments[0] + rewritten, unresolved = _rewrite_citation_markers_partial( + seg, + all_url_citations, + ) + if unresolved and not force: + pending_citation_segments[0] = rewritten + break + if unresolved and force: + rewritten = _replace_openai_citation_markers( + rewritten, + all_url_citations, + ) + pending_citation_segments.pop(0) + if rewritten: + out.append(rewritten) + return "".join(out) + + def _flush_pending_marker_tail(tail: str) -> str: + """Render any leftover citation tail at end-of-stream. + + Unterminated tails drop (no annotation to bind to). If the + close byte arrived concatenated, rewrite then scrub any + residual private-use bytes and any orphan ``cite`` + literal so the renderer never sees raw markup. url_citations + are aggregated separately and applied to web_search tool_end. + """ + if not tail: + return "" + if _OPENAI_CITE_STOP not in tail: + # Unterminated: drop the whole tail, otherwise the + # residual ``cite`` would leak as plain text. + return "" + rendered = _replace_openai_citation_markers( + tail, all_url_citations + ) + # Scrub residual private-use bytes (e.g. a partial opener). + for ch in ("", "", ""): + rendered = rendered.replace(ch, "") + # Drop any orphan ``cite`` literal -- meaningless + # without its closing byte and matching url_citation. + rendered = re.sub(r"^cite\S*", "", rendered) + return rendered def _emit_tool_event(payload: dict[str, Any]) -> str: + _stamp_server_tool_marker(payload) chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -3000,16 +5690,35 @@ class ExternalProviderClient: def _record_url_citation(payload: dict[str, Any]) -> None: """Append a url_citation onto the shared all_url_citations - list. Dedup by URL — the same source can be cited multiple - times across deltas. We do NOT try to attribute citations - to individual web_search_call invocations because OpenAI's - annotation events don't carry that linkage.""" + list. Dedup by URL — the same URL can be cited many + times under different ``source_id`` aliases (one per + span/locator), so collect every alias we see onto + the matching entry's ``source_ids`` list. The + delta-text rewriter resolves any of those aliases + back to this entry's URL. The id may live under + ``source_id``, ``id``, or ``locator`` across the + Responses API revisions.""" if payload.get("type") != "url_citation": return url = payload.get("url", "") if not url: return - if any(c["url"] == url for c in all_url_citations): + source_id = ( + payload.get("source_id") + or payload.get("id") + or payload.get("locator") + or "" + ) + # Single pass: either backfill aliases onto an + # existing URL entry (and return) or fall through + # to append a fresh one. + for c in all_url_citations: + if c["url"] != url: + continue + if source_id: + aliases = c.setdefault("source_ids", []) + if source_id not in aliases: + aliases.append(source_id) return title = payload.get("title") or url snippet = payload.get("snippet") or payload.get("quote") or "" @@ -3018,9 +5727,84 @@ class ExternalProviderClient: "url": url, "title": title, "snippet": snippet, + "source_ids": [source_id] if source_id else [], } ) + def _record_openai_reasoning_replay_item( + payload: Any, + ) -> Optional[dict[str, Any]]: + if not isinstance(payload, dict): + return None + item_id = payload.get("id") or payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + return None + existing = openai_reasoning_replay_items.setdefault( + item_id, + { + "type": "reasoning", + "id": item_id, + "summary": [], + "status": "completed", + }, + ) + if payload.get("type") == "reasoning": + sanitized = _sanitize_openai_reasoning_replay_item(payload) + if sanitized: + existing.update(sanitized) + return existing + summary_text = "" + part = payload.get("part") + if ( + isinstance(part, dict) + and part.get("type") == "summary_text" + ): + text = part.get("text") + if isinstance(text, str): + summary_text = text + elif ( + payload.get("type") + == "response.reasoning_summary_text.done" + ): + text = payload.get("text") + if isinstance(text, str): + summary_text = text + if summary_text: + summary_index = payload.get("summary_index") + summary = existing.setdefault("summary", []) + if isinstance(summary, list): + summary_part = { + "type": "summary_text", + "text": summary_text, + } + if ( + isinstance(summary_index, int) + and summary_index >= 0 + ): + while len(summary) <= summary_index: + summary.append( + {"type": "summary_text", "text": ""} + ) + summary[summary_index] = summary_part + else: + summary.append(summary_part) + return existing + + def _image_generation_arguments( + prompt: str, + raw_item_id: Any, + ) -> dict[str, Any]: + arguments: dict[str, Any] = {"kind": "image", "prompt": prompt} + if isinstance(raw_item_id, str) and raw_item_id: + arguments["openai_image_generation_call_id"] = raw_item_id + if current_openai_response_id: + arguments["openai_response_id"] = current_openai_response_id + if last_openai_reasoning_replay_item: + arguments["openai_reasoning_item"] = ( + last_openai_reasoning_replay_item + ) + return arguments + def _extract_reasoning_text(payload: Any) -> str: if payload is None: return "" @@ -3074,6 +5858,28 @@ class ExternalProviderClient: if not data_str: continue if data_str == "[DONE]": + # Flush any held-over partial marker; strip + # private-use bytes so garbled glyphs don't leak. + if pending_marker_tail: + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) + pending_marker_tail = "" + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) + # Force-drain any segment still awaiting an + # annotation; lingering codepoints are stripped. + tail_flushed = _drain_pending_segments( + force = True, + ) + if tail_flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(tail_flushed) if not done_emitted: yield "data: [DONE]" done_emitted = True @@ -3085,32 +5891,63 @@ class ExternalProviderClient: continue event_type = event.get("type") + _record_openai_response_id(event) if event_type == "response.output_text.delta": delta_text = event.get("delta", "") - if delta_text: - if reasoning_open: - yield _chunk_with_text("") - reasoning_open = False - yield _chunk_with_text(delta_text) - # Some API versions inline url citations on the - # delta event itself rather than as a separate - # response.output_text.annotation.added event. + # Process inline annotations first so source_ids + # referenced by same-delta markers are in the lookup + # before the rewriter runs. Some API versions inline + # url citations on the delta event itself. for ann in event.get("annotations") or []: if isinstance(ann, dict): _record_url_citation(ann) + if delta_text or pending_marker_tail: + # Prepend any held-over tail so a marker + # straddling two SSE events resolves cleanly. + combined = pending_marker_tail + delta_text + head, pending_marker_tail = ( + _split_pending_citation_tail(combined) + ) + if head: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + # Re-attempt earlier deferred segments first + # so output stays in order; the needed + # annotation may have arrived inline above. + flushed = _drain_pending_segments( + force = False, + ) + if flushed: + yield _chunk_with_text(flushed) + head_rewritten, has_unresolved = ( + _rewrite_citation_markers_partial( + head, + all_url_citations, + ) + ) + if has_unresolved or pending_citation_segments: + pending_citation_segments.append( + head_rewritten + ) + elif head_rewritten: + yield _chunk_with_text(head_rewritten) elif event_type == "response.output_text.annotation.added": ann = event.get("annotation") if isinstance(ann, dict): _record_url_citation(ann) + flushed = _drain_pending_segments( + force = False, + ) + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) elif event_type == "response.output_item.added": - # Track the call early but do NOT emit tool_start - # yet — action.query is not reliably populated on - # added across OpenAI API versions, and the - # frontend's tool_start is a one-shot push (no - # update mechanism). Wait for output_item.done. item = event.get("item", {}) if ( isinstance(item, dict) @@ -3120,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" @@ -3151,12 +5980,34 @@ class ExternalProviderClient: and latched_container_id is None ): latched_container_id = probe + if ( + isinstance(item, dict) + and item.get("type") == "image_generation_call" + ): + raw_item_id = item.get("id") + if isinstance(raw_item_id, str) and raw_item_id: + arguments = _image_generation_arguments( + "", + raw_item_id, + ) + image_generation_calls_started.add(raw_item_id) + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "image_generation", + "tool_call_id": raw_item_id, + "arguments": arguments, + } + ) elif event_type == "response.output_item.done": item = event.get("item", {}) if not isinstance(item, dict): continue if item.get("type") == "reasoning": + last_openai_reasoning_replay_item = ( + _record_openai_reasoning_replay_item(item) + ) summary_text = _extract_reasoning_text( item.get("summary") ) @@ -3196,14 +6047,16 @@ class ExternalProviderClient: ), } ) + # Per-card text; last call gets overwritten + # with citations at response.completed. + per_call_result = ( + f"Searching: {query}" if query else "" + ) yield _emit_tool_event( { "type": "tool_end", "tool_call_id": item_id, - # Empty result — the last call gets - # overwritten with citations at - # response.completed. - "result": "", + "result": per_call_result, } ) elif item.get("type") == "shell_call": @@ -3231,7 +6084,11 @@ class ExternalProviderClient: ) shell_calls.setdefault( item_id, - {"commands": [], "output": None}, + { + "commands": [], + "output": None, + "tool_end_emitted": False, + }, ) shell_calls[item_id]["commands"] = ( list(commands) @@ -3249,6 +6106,24 @@ class ExternalProviderClient: }, } ) + # Fallback: output may be bundled on the + # shell_call done event itself. + embedded_output = item.get("output") + if ( + isinstance(embedded_output, list) + and embedded_output + ): + shell_calls[item_id]["output"] = embedded_output + shell_calls[item_id]["tool_end_emitted"] = True + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": item_id, + "result": _format_shell_output( + embedded_output + ), + } + ) elif item.get("type") == "shell_call_output": # `call_id` links back to the shell_call's # `id`, which is what we used as the @@ -3259,8 +6134,15 @@ class ExternalProviderClient: item.get("call_id") or item.get("id") or "" ) output = item.get("output") or [] + # Skip if bundled-output path already + # finalised this card. + if shell_calls.get(call_id, {}).get( + "tool_end_emitted" + ): + continue if call_id in shell_calls: shell_calls[call_id]["output"] = output + shell_calls[call_id]["tool_end_emitted"] = True result_text = _format_shell_output(output) yield _emit_tool_event( { @@ -3270,43 +6152,29 @@ 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. - item_id = item.get("id", "") or ( - f"img_{time.time_ns()}" - ) + # 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 = ( item.get("revised_prompt") or item.get("prompt") or "" ) - yield _emit_tool_event( - { - "type": "tool_start", - "tool_name": "image_generation", - "tool_call_id": item_id, - "arguments": { - "kind": "image", - "prompt": prompt_in, - }, - } + done_arguments = _image_generation_arguments( + prompt_in, + raw_item_id, ) + if item_id not in image_generation_calls_started: + yield _emit_tool_event( + { + "type": "tool_start", + "tool_name": "image_generation", + "tool_call_id": item_id, + "arguments": done_arguments, + } + ) b64 = ( item.get("result") or item.get("b64_json") or "" ) @@ -3316,18 +6184,75 @@ class ExternalProviderClient: "type": "tool_end", "tool_call_id": item_id, "result": "", + "arguments": done_arguments, "image_b64": b64, "image_mime": (f"image/{output_format}"), "size": item.get("size"), "quality": item.get("quality"), "background": item.get("background"), + "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) and "reasoning" in event_type ): + recorded_reasoning = ( + _record_openai_reasoning_replay_item(event) + ) + if recorded_reasoning: + last_openai_reasoning_replay_item = ( + recorded_reasoning + ) reasoning_delta = _extract_reasoning_text(event) if reasoning_delta: if not reasoning_open: @@ -3342,6 +6267,34 @@ class ExternalProviderClient: ) if isinstance(completed_usage, dict): last_usage = completed_usage + # Flush any unterminated citation tail + # held over from the last delta. By + # the time we get here every annotation + # has been recorded so a late-arriving + # source_id may resolve cleanly; if it + # still doesn't, the helper strips the + # private-use bytes so no garbled + # glyph reaches the user. + if pending_marker_tail: + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) + pending_marker_tail = "" + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) + # Force-drain any segment still awaiting an + # annotation; lingering codepoints are stripped. + tail_flushed = _drain_pending_segments( + force = True, + ) + if tail_flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(tail_flushed) if reasoning_open: yield _chunk_with_text("") reasoning_open = False @@ -3378,22 +6331,16 @@ class ExternalProviderClient: } ) container_id_emitted = True - # Apply the aggregated citation list onto the - # *last* web_search call by overwriting its - # tool_end result. The frontend's - # parseSourcesFromResult flatMaps every - # web_search tool-call result, so a single - # non-empty result is enough to surface the - # whole source-pill set at the message tail — - # no need to fan out across every card (which - # would just duplicate the same pills). + # Overwrite the last web_search call with the + # citation list; the source-pill extractor + # flatMaps across cards. Earlier cards keep + # their per-call "Searching:" text. if web_search_calls and all_url_citations: last_id = list(web_search_calls.keys())[-1] blocks: list[str] = [] for cit in all_url_citations: line = ( - f"Title: {cit['title']}\n" - f"URL: {cit['url']}" + f"Title: {cit['title']}\nURL: {cit['url']}" ) if cit.get("snippet"): line += f"\nSnippet: {cit['snippet']}" @@ -3405,6 +6352,21 @@ class ExternalProviderClient: "result": "\n---\n".join(blocks), } ) + # Final flush: finalise any orphan shell_call + # so the card stops spinning. + for sc_id, sc_state in shell_calls.items(): + if sc_state.get("tool_end_emitted"): + continue + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": sc_id, + "result": _format_shell_output( + sc_state.get("output") or [] + ), + } + ) + sc_state["tool_end_emitted"] = True chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -3412,7 +6374,11 @@ class ExternalProviderClient: { "index": 0, "delta": {}, - "finish_reason": "stop", + "finish_reason": ( + "tool_calls" + if saw_function_call + else "stop" + ), } ], } @@ -3434,6 +6400,29 @@ class ExternalProviderClient: ) if isinstance(incomplete_usage, dict): last_usage = incomplete_usage + # Same flush as response.completed -- + # truncated streams can leave a half- + # marker in the buffer. + if pending_marker_tail: + flushed = _flush_pending_marker_tail( + pending_marker_tail + ) + pending_marker_tail = "" + if flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(flushed) + # Force-drain any segment still awaiting an + # annotation; lingering codepoints are stripped. + tail_flushed = _drain_pending_segments( + force = True, + ) + if tail_flushed: + if reasoning_open: + yield _chunk_with_text("") + reasoning_open = False + yield _chunk_with_text(tail_flushed) if reasoning_open: yield _chunk_with_text("") reasoning_open = False @@ -3448,8 +6437,7 @@ class ExternalProviderClient: blocks = [] for cit in all_url_citations: line = ( - f"Title: {cit['title']}\n" - f"URL: {cit['url']}" + f"Title: {cit['title']}\nURL: {cit['url']}" ) if cit.get("snippet"): line += f"\nSnippet: {cit['snippet']}" @@ -3461,6 +6449,22 @@ class ExternalProviderClient: "result": "\n---\n".join(blocks), } ) + # Mirror the response.completed flush so + # truncated streams also finalise orphan + # shell_calls. + for sc_id, sc_state in shell_calls.items(): + if sc_state.get("tool_end_emitted"): + continue + yield _emit_tool_event( + { + "type": "tool_end", + "tool_call_id": sc_id, + "result": _format_shell_output( + sc_state.get("output") or [] + ), + } + ) + sc_state["tool_end_emitted"] = True chunk = { "id": completion_id, "object": "chat.completion.chunk", @@ -3655,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("/") @@ -3939,6 +7002,17 @@ def _build_usage_chunk( "cache_creation_input_tokens": cache_creation, "cache_read_input_tokens": cache_read, } + # Forward 5m/1h cache-write breakdown so cost calc applies the + # 2x 1h premium instead of defaulting to 5m on chat-style. + cc_breakdown = last_usage.get("cache_creation") + if isinstance(cc_breakdown, dict) and cc_breakdown: + usage_block["cache_creation"] = cc_breakdown + # Propagate fast-mode `usage.speed` so the cost ledger can apply + # the 6x multiplier without re-derivation (Anthropic falls back + # to "standard" when fast-mode is unsupported or rate-limited). + speed = last_usage.get("speed") + if speed in ("fast", "standard"): + usage_block["speed"] = speed else: prompt_tokens = last_usage.get("input_tokens") or 0 cached = 0 @@ -3953,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 b8d168ccc5..4acc2310e3 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -60,7 +60,9 @@ _INTENT_SIGNAL = re.compile( # Handles both straight and curly apostrophes. # Excludes "I can", "I should", "I want to", "let's" which # appear frequently in direct answers / explanations. - r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b" + # Negative lookahead drops negated forms ("I will not", "I'll never") + # so a refusal doesn't trigger a re-prompt. + r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" r"|" # Step/plan framing: "First ...", "Step 1:", "Here's my plan" r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" @@ -5091,14 +5093,31 @@ 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, - tool_context = tool_context, - ) + # 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, + tool_context = tool_context, + ) 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..4f528a689f 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] = [] @@ -124,15 +114,15 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: 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,9 +159,8 @@ _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"} ) @@ -187,14 +176,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 +202,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..a5e614899d --- /dev/null +++ b/studio/backend/core/inference/mcp_client.py @@ -0,0 +1,181 @@ +# 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 +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +MCP_TOOL_PREFIX = "mcp__" + +_oauth_token_store = None + + +def parse_server_headers(server: dict) -> Optional[dict]: + 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 + 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/pricing.py b/studio/backend/core/inference/pricing.py index 74c57fa594..84eec17e84 100644 --- a/studio/backend/core/inference/pricing.py +++ b/studio/backend/core/inference/pricing.py @@ -1,50 +1,24 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Static per-MTok pricing tables for external providers, plus a -``calculate_cost`` helper that turns an upstream ``usage`` block into -a USD figure for surfacing in the chat UI. +"""Static per-MTok pricing tables and ``calculate_cost`` helper for +turning an upstream ``usage`` block into a USD figure. -Neither the Anthropic Messages API nor the OpenAI Responses API -reports a ``cost`` field on the response. Both expose detailed token -counts (input, output, cache hits, server-tool invocations); pricing -multipliers live in the provider docs. We fold the docs into a static -table here, multiply by the usage block, and emit a per-turn cost + -running session total client-side. - -Sources (verified live 2026-05-22): -- Anthropic models overview: - https://platform.claude.com/docs/en/about-claude/models/overview -- Anthropic prompt-caching multipliers (5m write 1.25x, 1h write 2x, - read 0.1x): - https://platform.claude.com/docs/en/build-with-claude/prompt-caching -- Anthropic web search ($10 / 1000 searches, code execution - free-with-paid when paired with the newer web tools): - https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool - https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool -- OpenAI pricing page (input / output per MTok per model family): - https://platform.openai.com/docs/pricing +Sources: Anthropic prompt-caching docs (5m write 1.25x, 1h write 2x, +read 0.1x), web search ($10/1000), code execution; OpenAI pricing page. """ from __future__ import annotations from typing import Any, Optional -# Per-million-token base pricing. `cache_5m_write_mult`, `cache_1h_write_mult`, -# `cache_read_mult` are multipliers ON `input_per_mtok` -- not absolute prices -- -# matching how Anthropic publishes them (5m write = 1.25x base, etc.). -# -# `input_per_mtok` and `output_per_mtok` are USD per 1,000,000 tokens. +# Per-MTok base pricing in USD. Cache multipliers are applied ON +# `input_per_mtok` (not absolute prices), matching Anthropic's docs. ANTHROPIC_PRICING: dict[str, dict[str, float]] = { "claude-opus-4-7": {"input_per_mtok": 5.0, "output_per_mtok": 25.0}, "claude-opus-4-6": {"input_per_mtok": 5.0, "output_per_mtok": 25.0}, - # Canonical 4.5 ids are referenced from backend defaults (e.g. - # PROVIDER_REGISTRY['anthropic'].default_models) without the date - # suffix. The dated ids ARE the canonical names per Anthropic's - # models overview, but lookups for the bare id ("claude-opus-4-5") - # don't prefix-match the dated key the other way around, so we - # alias both forms here. Otherwise calculate_cost returns - # priced=False + zero cost for the common ids. + # Alias both the bare id and dated id: backend defaults reference + # the bare form, which won't prefix-match the dated key. "claude-opus-4-5": {"input_per_mtok": 5.0, "output_per_mtok": 25.0}, "claude-opus-4-5-20251101": {"input_per_mtok": 5.0, "output_per_mtok": 25.0}, "claude-opus-4-1": {"input_per_mtok": 15.0, "output_per_mtok": 75.0}, @@ -59,19 +33,9 @@ ANTHROPIC_PRICING: dict[str, dict[str, float]] = { } OPENAI_PRICING: dict[str, dict[str, float]] = { - # All values verified against developers.openai.com/api/docs/pricing - # 2026-05-22. Update against the live pricing page on every model launch. - # Initial commit underbilled every gpt-5.x family 2-6x -- fixed here - # after PR review caught it via doc cross-check. - # - # `long_context_input_per_mtok` / `long_context_output_per_mtok` / - # `long_context_threshold` are populated when OpenAI publishes a - # second pricing tier for prompts above N input tokens. gpt-5.5 and - # gpt-5.4 cross over at 272k input tokens; the long-context rates - # are double the headline input price (and ~1.5x on output). Other - # families currently ship with a single rate (no `long_context_*` - # keys = no tier crossover). Reference: - # https://developers.openai.com/api/docs/pricing + # Verified against developers.openai.com/api/docs/pricing. + # `long_context_*` keys apply once input exceeds the threshold + # (gpt-5.5/5.4: 272k); families without these keys ship a single rate. "gpt-5.5": { "input_per_mtok": 5.0, "output_per_mtok": 30.0, @@ -91,43 +55,33 @@ OPENAI_PRICING: dict[str, dict[str, float]] = { "gpt-5.4-mini": {"input_per_mtok": 0.75, "output_per_mtok": 4.5}, "gpt-5.4-nano": {"input_per_mtok": 0.20, "output_per_mtok": 1.25}, "gpt-5.3-codex": {"input_per_mtok": 1.75, "output_per_mtok": 14.0}, - # chat-latest / gpt-5.3-chat-latest is an alias for the current - # ChatGPT model; same price as gpt-5.5. + # chat-latest aliases gpt-5.5. "gpt-5.3-chat-latest": {"input_per_mtok": 5.0, "output_per_mtok": 30.0}, "chat-latest": {"input_per_mtok": 5.0, "output_per_mtok": 30.0}, - # o-series and gpt-4.5: NOT currently listed on the pricing page. - # Removed to avoid silent-underbilling drift. Returning priced=False - # is honest; the UI can still render token counts. Restore with - # verified per-MTok rates if/when the page lists them again. + # o-series and gpt-4.5 are no longer on the pricing page; omit them + # so calculate_cost returns priced=False rather than silently $0. } # Shared multipliers (same across every Anthropic model). ANTHROPIC_CACHE_5M_WRITE_MULT = 1.25 ANTHROPIC_CACHE_1H_WRITE_MULT = 2.0 ANTHROPIC_CACHE_READ_MULT = 0.1 +# Anthropic fast-mode (Opus 4.6 / 4.7 only): 6x standard on input + output. +# https://platform.claude.com/docs/en/build-with-claude/fast-mode#pricing +ANTHROPIC_FAST_MODE_MULT = 6.0 -# OpenAI: cache reads are 0.1x base input, cache writes are not billed -# separately (the first prefix-write request just pays normal input). +# OpenAI: cache reads 0.1x; cache writes pay normal input price. OPENAI_CACHE_READ_MULT = 0.1 -# Server-tool surcharges. -# Anthropic: $10 / 1000 web searches; code_execution is $0.05/hr after -# 50 free hours/day per org (no per-org visibility here, so the -# calculator reports the marginal rate). +# Server-tool surcharges. Anthropic code_exec is $0.05/hr marginal +# (50 free hours/day per org, not visible here). ANTHROPIC_WEB_SEARCH_USD_PER_1K = 10.0 ANTHROPIC_CODE_EXEC_USD_PER_HOUR = 0.05 -# OpenAI: web_search is billed at $10/1000 calls plus the model's -# token rate for the returned search content (already captured under -# input/output_tokens). The hosted shell tool bills per 20-minute -# session per container memory tier (1g/4g/16g/64g at -# $0.03/$0.12/$0.48/$1.92). Since Studio doesn't surface the memory -# tier in the cost ledger and most users land on the default 1g, we -# bill the 1g rate ($0.09/hour) and let the user inspect the OpenAI -# dashboard for the exact figure on heavier configs. -# Source: developers.openai.com/api/docs/pricing 2026-05-22. +# OpenAI container bills per memory tier; we report the 1g default +# ($0.09/hour) since the tier isn't surfaced to the cost ledger. OPENAI_WEB_SEARCH_USD_PER_1K = 10.0 -OPENAI_CONTAINER_USD_PER_HOUR = 0.09 # 1g default tier; 3 x $0.03 / 60min +OPENAI_CONTAINER_USD_PER_HOUR = 0.09 # 1g default tier def _lookup(provider: str, model: str) -> Optional[dict[str, float]]: @@ -142,11 +96,13 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]: return None if model in table: return table[model] - # Fall back to a prefix match so date-suffixed snapshots - # ("gpt-5.5-2026-04-23") inherit the canonical-id prices. - for key, val in table.items(): - if model.startswith(key): - return val + # Longest-prefix match on a dash boundary: lets dated snapshots + # inherit canonical prices while preventing "claude-opus-4-15" + # from matching "claude-opus-4-1" or "gpt-5.5-prod" from matching + # "gpt-5.5-pro". Sort longest-first to pick the most specific row. + for key in sorted(table, key = len, reverse = True): + if model.startswith(key) and (len(model) == len(key) or model[len(key)] == "-"): + return table[key] return None @@ -155,28 +111,11 @@ def calculate_cost( model: str, usage: dict[str, Any], ) -> dict[str, float]: - """Return a per-turn USD cost breakdown. - - Returns a dict with the per-bucket cost AND the totals so the - frontend can render either a single number or a "where did the - money go" tooltip without re-doing the math: - - { - "input_usd": 0.0042, - "output_usd": 0.012, - "cache_write_usd": 0.0001, - "cache_read_usd": 0.0008, - "server_tools_usd": 0.01, - "total_usd": 0.0271, - "billable_input_tokens": 5023, # input + cache_create + cache_read - "billable_output_tokens": 480, - "model_priced": "claude-opus-4-7", - "priced": true, - } - - When the model isn't in the static table (new family, custom base - URL), `priced` is False and every USD field is 0.0; the frontend - can still show the token counts. + """Return a per-turn USD cost breakdown with per-bucket + total + fields so the frontend can render either a single number or a + tooltip without re-doing the math. When the model isn't in the + static table, ``priced`` is False and USD fields are 0.0 (token + counts still report). """ prices = _lookup(provider, model) out: dict[str, float] = { @@ -192,34 +131,64 @@ def calculate_cost( "priced": bool(prices), } - input_tokens = int(usage.get("input_tokens") or 0) - output_tokens = int(usage.get("output_tokens") or 0) - cache_creation = int(usage.get("cache_creation_input_tokens") or 0) - cache_read = int(usage.get("cache_read_input_tokens") or 0) - # OpenAI Responses reports cached tokens under input_tokens_details - # but ALSO folds them into the top-level input_tokens, so we don't - # add cache_read into the billable total again below (Anthropic - # excludes cache buckets from input_tokens, OpenAI includes them -- - # the two providers differ here and the calculator must match). - if provider == "openai": - details = usage.get("input_tokens_details") or {} + # Accept raw (input_tokens/output_tokens) and Studio chat-style + # (prompt_tokens/completion_tokens) envelopes. Cache buckets + # behave differently per envelope: + # raw Anthropic: input_tokens EXCLUDES cache buckets + # raw OpenAI: input_tokens INCLUDES cache_read + # Studio Anthropic: prompt_tokens INCLUDES cache_creation + cache_read + # Studio OpenAI: prompt_tokens == raw input_tokens + # Clamp tokens >=0 so corrupted payloads can't produce a negative bill. + cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0)) + cache_read_native_present = ( + "cache_read_input_tokens" in usage + and usage.get("cache_read_input_tokens") is not None + ) + cache_read = max(0, int(usage.get("cache_read_input_tokens") or 0)) + # Fallback to mirrored prompt_tokens_details only when the native + # cache_read_input_tokens key is absent. An explicit native 0 is + # authoritative, so a stale mirrored block from a proxy can never + # inflate cache_read past the native count. + if not cache_read_native_present: + details = usage.get("prompt_tokens_details") or {} if isinstance(details, dict): - cache_read = max(cache_read, int(details.get("cached_tokens") or 0)) - # OpenAI: cache_read already counted inside input_tokens. + cache_read = max(0, int(details.get("cached_tokens") or 0)) + has_input_tokens = "input_tokens" in usage and usage.get("input_tokens") is not None + if has_input_tokens: + input_tokens = max(0, int(usage.get("input_tokens") or 0)) + else: + # Chat-style: peel cache buckets back out for Anthropic to + # recover the raw uncached prompt count. + prompt_tokens = max(0, int(usage.get("prompt_tokens") or 0)) + if provider == "anthropic": + input_tokens = max(0, prompt_tokens - cache_creation - cache_read) + else: + input_tokens = prompt_tokens + # Prefer raw output_tokens even when 0 (an `or` fallback would + # silently pick a stale completion_tokens). + if "output_tokens" in usage and usage.get("output_tokens") is not None: + output_tokens = max(0, int(usage.get("output_tokens") or 0)) + else: + output_tokens = max(0, int(usage.get("completion_tokens") or 0)) + if provider == "openai": + # Cached tokens land on either input_tokens_details (raw + # Responses) or prompt_tokens_details (Studio chat-style). + for key in ("input_tokens_details", "prompt_tokens_details"): + details = usage.get(key) or {} + if isinstance(details, dict): + cache_read = max(cache_read, int(details.get("cached_tokens") or 0)) + # OpenAI input_tokens already counts cache_read. out["billable_input_tokens"] = input_tokens + cache_creation else: - # Anthropic: input_tokens excludes cache_* buckets, add them all. + # Anthropic input_tokens excludes cache buckets; add them back. out["billable_input_tokens"] = input_tokens + cache_creation + cache_read out["billable_output_tokens"] = output_tokens if not prices: return out - # Long-context tier crossover (gpt-5.5 / gpt-5.4 today). OpenAI - # bills the whole turn at the long-context rate once the prompt - # crosses the threshold, NOT a per-token blend, so we pick a - # single (base, out_per) pair for this turn based on - # billable_input_tokens. + # Long-context tier: whole-turn flip (not per-token blend) once + # billable_input_tokens crosses the threshold. lc_thresh = prices.get("long_context_threshold") in_long_context_tier = ( lc_thresh is not None @@ -235,17 +204,27 @@ def calculate_cost( base = prices["input_per_mtok"] out_per = prices["output_per_mtok"] + # Anthropic fast-mode: 6x on input + output. Cache multipliers stack + # on top of fast-mode, so applying once to (base, out_per) propagates + # into the cache_*_usd buckets computed below. + if provider == "anthropic" and usage.get("speed") == "fast": + base *= ANTHROPIC_FAST_MODE_MULT + out_per *= ANTHROPIC_FAST_MODE_MULT + if out["model_priced"]: + out["model_priced"] = f"{out['model_priced']} (fast)" + out["input_usd"] = (input_tokens / 1_000_000.0) * base out["output_usd"] = (output_tokens / 1_000_000.0) * out_per if provider == "anthropic": - # Split cache_creation across 5m / 1h buckets when the - # response surfaces the breakdown. - cc_breakdown = usage.get("cache_creation") or {} - cc_5m = int(cc_breakdown.get("ephemeral_5m_input_tokens") or 0) - cc_1h = int(cc_breakdown.get("ephemeral_1h_input_tokens") or 0) + # Split cache_creation into 5m / 1h buckets when surfaced. + # Tolerate non-dict (some proxies fold to an int total). + cc_raw = usage.get("cache_creation") + cc_breakdown = cc_raw if isinstance(cc_raw, dict) else {} + cc_5m = max(0, int(cc_breakdown.get("ephemeral_5m_input_tokens") or 0)) + cc_1h = max(0, int(cc_breakdown.get("ephemeral_1h_input_tokens") or 0)) if cc_5m + cc_1h == 0 and cache_creation > 0: - # Fall back: assume default 5m pool when no breakdown is given. + # No breakdown -- assume default 5m pool. cc_5m = cache_creation out["cache_write_usd"] = ( cc_5m / 1_000_000.0 @@ -265,24 +244,17 @@ def calculate_cost( + code_exec_hours * ANTHROPIC_CODE_EXEC_USD_PER_HOUR ) else: - # OpenAI: cache writes share the base input price (no premium). - # Only cache reads get the 0.1x multiplier; subtract those from - # the input_usd we already counted so we don't double-bill. - # Anthropic excludes cache buckets from input_tokens, but - # OpenAI folds them in, so the math differs. + # OpenAI: cache writes pay base input; only cache reads get + # 0.1x. Subtract cached from already-counted input_usd to + # avoid double-billing (OpenAI folds cache into input_tokens). if cache_read > 0: non_cached_input = max(0, input_tokens - cache_read) out["input_usd"] = (non_cached_input / 1_000_000.0) * base out["cache_read_usd"] = ( (cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT ) - # Server-tool surcharges. OpenAI doesn't include these on its - # `usage` object directly -- web_search invocations are counted - # from `ResponseFunctionWebSearch` items in the output array, - # and container hours come from the SSE translator's shell-tool - # accounting. Studio surfaces both under a normalised - # `openai_tool_use` key on the usage dict the SSE finaliser - # hands to this calculator. + # OpenAI server-tool surcharges arrive under `openai_tool_use` + # (normalised by the SSE finaliser from output array items). srv = usage.get("openai_tool_use") or {} if isinstance(srv, dict): web_searches = int(srv.get("web_search_requests") or 0) @@ -304,17 +276,14 @@ def calculate_cost( def pricing_snapshot() -> dict[str, Any]: - """Whole pricing table, for the /api/providers/pricing endpoint. - - Returns a flat structure the frontend can hand to its cost - formatter without re-implementing the multipliers. - """ + """Whole pricing table for the /api/providers/pricing endpoint.""" return { "anthropic": { "models": dict(ANTHROPIC_PRICING), "cache_5m_write_mult": ANTHROPIC_CACHE_5M_WRITE_MULT, "cache_1h_write_mult": ANTHROPIC_CACHE_1H_WRITE_MULT, "cache_read_mult": ANTHROPIC_CACHE_READ_MULT, + "fast_mode_mult": ANTHROPIC_FAST_MODE_MULT, "web_search_usd_per_1k": ANTHROPIC_WEB_SEARCH_USD_PER_1K, "code_execution_usd_per_hour": ANTHROPIC_CODE_EXEC_USD_PER_HOUR, }, 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 21ecb9edf4..48d70aa67f 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,14 @@ import tempfile import threading import urllib.request +from core.inference.mcp_client import ( + MCP_TOOL_PREFIX, + call_tool_sync, + list_tools_async, + parse_server_headers, +) +from storage import mcp_servers_db + from loggers import get_logger logger = get_logger(__name__) @@ -513,6 +522,92 @@ def _get_rag_tool_spec(): ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL, _get_rag_tool_spec()] +# 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")] + if not servers: + return [] + + # OAuth probes need minutes for first-connect/expired-token browser + # sign-in; non-OAuth probes fail fast. Matches routes/mcp_servers.py. + results = await asyncio.gather( + *( + list_tools_async( + url = s["url"], + headers = parse_server_headers(s), + timeout = 305.0 if s.get("use_oauth") else 8.0, + 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() @@ -538,6 +633,25 @@ 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" + 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", ""), @@ -662,8 +776,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..b9643cac6a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3057,6 +3057,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 +3131,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..632b38d75a 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -959,7 +959,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 +1019,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 +1030,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 @@ -1168,6 +1208,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 +1361,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 +1379,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})...") @@ -2248,6 +2313,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, diff --git a/studio/backend/main.py b/studio/backend/main.py index d1a61d4417..1d60bba881 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -18,6 +18,17 @@ _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 @@ -49,6 +60,8 @@ import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version +from typing import Optional +from urllib.parse import urlparse _STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$") @@ -120,6 +133,7 @@ from routes import ( export_router, inference_router, inference_studio_router, + mcp_servers_router, models_router, providers_router, rag_router, @@ -523,6 +537,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"]) @@ -717,10 +732,8 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes: def _inject_bootstrap(html_bytes: bytes, app: FastAPI): """Inject bootstrap credentials when password change is pending. - - Returns ``(html_bytes, script_nonce_or_None)``. Callers must forward - the nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so the inline script is - not blocked by CSP. + Returns ``(html_bytes, script_nonce_or_None)``; callers forward the + nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so CSP allows the inline script. """ import json as _json import secrets as _secrets @@ -745,6 +758,86 @@ def _inject_bootstrap(html_bytes: bytes, app: FastAPI): return html.encode("utf-8"), nonce +_DEFAULT_PORTS = {"http": 80, "https": 443, "ws": 80, "wss": 443} + + +def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]]: + """Canonicalise an Origin to ``(scheme, host, port)`` for equality. + Browsers strip default ports (RFC 6454 sec 6.1) and scheme/host are + case-insensitive (RFC 3986), so bare string compare misclassifies + same-origin requests as cross-origin. Returns ``None`` on unparseable + input so callers fall to the safer cross-origin default. + """ + scheme = (scheme or "").strip().lower() + if not scheme or not netloc: + return None + # Strip userinfo (RFC 3986); Origin never carries credentials. + if "@" in netloc: + netloc = netloc.rsplit("@", 1)[1] + # IPv6 hosts use brackets (RFC 3986 sec 3.2.2): ``[::1]:8902``. Bare + # ``partition(":")`` mis-parses these and breaks ``unsloth studio -H ::1``. + if netloc.startswith("["): + close = netloc.find("]") + if close == -1: + return None + host = netloc[1:close] + rest = netloc[close + 1 :] + if rest.startswith(":"): + port_str = rest[1:] + elif rest == "": + port_str = "" + else: + return None + else: + host, _, port_str = netloc.partition(":") + host = host.strip().lower() + if not host: + return None + if port_str: + try: + port = int(port_str) + except ValueError: + return None + else: + port = _DEFAULT_PORTS.get(scheme, 0) + return (scheme, host, port) + + +def _is_same_origin_request(request: Request) -> bool: + """True when Origin is missing or matches request's scheme://host:port. + Top-level same-document GETs omit Origin, so missing counts as same-origin. + Callers must also emit ``Vary: Origin``. Both sides are canonicalised via + :func:`_canonical_origin` so default-port stripping and scheme/host case + do not misclassify same-origin requests as cross-origin. + """ + origin = request.headers.get("origin") + if origin is None: + # Missing header: top-level same-document GETs omit Origin. + return True + # Empty string is not a valid serialised origin (RFC 6454 sec 6.1). + if not origin: + return False + # "null" token (sandboxed iframes, file:// pages) is never same-origin. + if origin == "null": + return False + # ``urlparse`` raises ``ValueError`` on malformed IPv6 brackets; swallow + # so a garbage Origin doesn't 500 the SPA handler. + try: + parsed = urlparse(origin) + except ValueError: + return False + origin_canon = _canonical_origin(parsed.scheme, parsed.netloc) + if origin_canon is None: + return False + try: + self_canon = _canonical_origin(request.url.scheme, request.url.netloc) + except ValueError: + return False + if self_canon is None: + return False + return origin_canon == self_canon + + def setup_frontend(app: FastAPI, build_path: Path): """Mount frontend static files (optional)""" if not build_path.exists(): @@ -755,11 +848,18 @@ def setup_frontend(app: FastAPI, build_path: Path): if assets_dir.exists(): app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") - def _build_index_response() -> Response: + def _build_index_response(request: Request) -> Response: content = (build_path / "index.html").read_bytes() content = _strip_crossorigin(content) - content, nonce = _inject_bootstrap(content, app) - headers = {"Cache-Control": "no-cache, no-store, must-revalidate"} + # Bootstrap pw is same-origin only; Vary: Origin keeps caches honest. + if _is_same_origin_request(request): + content, nonce = _inject_bootstrap(content, app) + else: + nonce = None + headers = { + "Cache-Control": "no-cache, no-store, must-revalidate", + "Vary": "Origin", + } if nonce: headers[_CSP_SCRIPT_NONCE_HEADER] = nonce return Response( @@ -769,11 +869,11 @@ def setup_frontend(app: FastAPI, build_path: Path): ) @app.get("/") - async def serve_root(): - return _build_index_response() + async def serve_root(request: Request): + return _build_index_response(request) @app.get("/{full_path:path}") - async def serve_frontend(full_path: str): + async def serve_frontend(request: Request, full_path: str): if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")): return {"error": "API endpoint not found"} @@ -787,6 +887,6 @@ def setup_frontend(app: FastAPI, build_path: Path): return FileResponse(file_path) # Serve index.html as bytes — avoids Content-Length mismatch - return _build_index_response() + return _build_index_response(request) return True diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b4d6bb926b..e585e4c346 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -299,7 +299,11 @@ class InferenceStatusResponse(BaseModel): """Current inference backend status""" active_model: Optional[str] = Field( - None, description = "Currently active model identifier" + None, description = "Currently active model display identifier" + ) + model_identifier: Optional[str] = Field( + None, + description = "Loadable identifier for the active model.", ) is_vision: bool = Field( False, description = "Whether the active model is a vision model" @@ -471,6 +475,40 @@ class InputDocumentContentPart(BaseModel): ) +class OpenAIReasoningContentPart(BaseModel): + """OpenAI Responses reasoning item paired with a tool output. + + Reasoning models can require the previous ``reasoning`` output item + to be replayed immediately before an ``image_generation_call`` id + when manually managing Responses context. This part is OpenAI-only; + routes strip it for every other provider before proxying. + """ + + type: Literal["reasoning"] + id: str = Field(..., description = "OpenAI reasoning output item id.") + summary: list[dict[str, Any]] = Field(default_factory = list) + status: Optional[Literal["in_progress", "completed", "incomplete"]] = None + + +class ImageGenerationCallContentPart(BaseModel): + """OpenAI Responses image_generation call reference. + + OpenAI accepts prior ``image_generation_call`` items in the next + Responses ``input`` array so follow-up prompts can edit or refine a + generated image without resending the base64 payload. The frontend + forwards this as a synthetic assistant content part when building + the next OpenAI Responses request; ``external_provider`` translates + it back to the provider-specific top-level input item. + """ + + type: Literal["image_generation_call"] + id: str = Field(..., description = "OpenAI image_generation_call output item id.") + response_id: Optional[str] = Field( + None, + description = "OpenAI Responses response id to use as previous_response_id for follow-up edits.", + ) + + class CompactionContentPart(BaseModel): """Anthropic server-side compaction state, attached to an assistant message for round-tripping on the next turn. @@ -504,6 +542,8 @@ ContentPart = Annotated[ Annotated[TextContentPart, Tag("text")], Annotated[ImageContentPart, Tag("image_url")], Annotated[InputDocumentContentPart, Tag("input_document")], + Annotated[OpenAIReasoningContentPart, Tag("reasoning")], + Annotated[ImageGenerationCallContentPart, Tag("image_generation_call")], Annotated[CompactionContentPart, Tag("compaction")], ], Discriminator(_content_part_discriminator), @@ -541,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": @@ -668,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.", @@ -723,17 +775,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 = ( @@ -797,6 +874,16 @@ class ChatCompletionRequest(BaseModel): "to auto-create." ), ) + fast_mode: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Anthropic fast-mode toggle. On Claude Opus 4.6 / " + "4.7 adds the `fast-mode-2026-02-01` beta header and sends " + "`speed: 'fast'` for higher OTPS at premium pricing. Silently " + "ignored on every other model + provider. See " + "https://platform.claude.com/docs/en/build-with-claude/fast-mode" + ), + ) @model_validator(mode = "after") def _resolve_missing_tool_call_ids(self) -> "ChatCompletionRequest": 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/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index c33ebf4d94..85294114b1 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -22,19 +22,11 @@ rich>=13.0 markdown-it-py>=3.0 mdurl>=0.1 pygments>=2.0 -# pydantic is intentionally NOT installed via this --no-deps file. -# install.sh / install.ps1 / install_python_stack.py run a separate -# `pip install pydantic` (with deps) just before this file is -# applied, so pip resolves `pydantic-core` to the exact version -# pydantic's `_ensure_pydantic_core_version` check expects. Listing -# pydantic + pydantic-core unpinned here and resolving them under -# --no-deps used to pick the latest of each independently and trip -# `SystemError: pydantic-core 2.X.Y is incompatible with the current -# pydantic version` on the first import (Windows fresh-venv repro -# was the canonical case). pydantic's transitive deps -# (annotated-types, pydantic-core, typing-extensions, -# typing-inspection) are torch-free, so installing it WITH deps -# does not pull torch. +# pydantic is intentionally NOT pinned here. install.sh / install.ps1 +# / install_python_stack.py run `pip install pydantic` WITH deps just +# before this --no-deps file is applied, so pip resolves pydantic-core +# to the exact version pydantic's _ensure_pydantic_core_version check +# expects. Pinning both under --no-deps used to drift them apart. pyyaml nest-asyncio diff --git a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt new file mode 100644 index 0000000000..2cd03d8b78 --- /dev/null +++ b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt @@ -0,0 +1,5 @@ +# mlx-vlm / mlx-lm declare transformers>=5.x which conflicts with the +# main venv's constraints.txt pin transformers==4.57.6 and forces uv to +# backtrack unsloth. Relax to match the pin -- per-model 5.x routing +# happens at runtime via the side-car venvs. +transformers>=4.57.6 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 433355f969..49cd05a3fd 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -17,6 +17,7 @@ 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.rag import router as rag_router +from routes.mcp_servers import router as mcp_servers_router __all__ = [ "training_router", @@ -31,4 +32,5 @@ __all__ = [ "chat_history_router", "providers_router", "rag_router", + "mcp_servers_router", ] diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index da6416e324..107a1657f3 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -95,6 +95,89 @@ def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]: return aliases +def _used_local_model_selections( + recipe: dict[str, Any], local_provider_names: set[str] +) -> dict[tuple[str, str], list[str]]: + used_aliases = _used_llm_model_aliases(recipe) + selections: dict[tuple[str, str], list[str]] = {} + for mc in recipe.get("model_configs", []): + if not isinstance(mc, dict): + continue + alias = mc.get("alias") + if not isinstance(alias, str) or alias not in used_aliases: + continue + provider = mc.get("provider") + if not isinstance(provider, str) or provider not in local_provider_names: + continue + model = mc.get("model") + target = model.strip() if isinstance(model, str) else "" + if not target or target.lower() == "local": + continue + variant = mc.get("gguf_variant") + gguf_variant = variant.strip() if isinstance(variant, str) else "" + selections.setdefault((target, gguf_variant), []).append(alias) + return selections + + +def _single_used_local_model_selection( + recipe: dict[str, Any], local_provider_names: set[str] +) -> tuple[str, str] | None: + selections = _used_local_model_selections(recipe, local_provider_names) + if not selections: + return None + if len(selections) > 1: + aliases = ", ".join(alias for values in selections.values() for alias in values) + raise ValueError( + "Recipes supports one active local model per run. " + f"Select the same local model and GGUF variant for: {aliases}." + ) + return next(iter(selections)) + + +def _loaded_local_model_identity() -> tuple[bool, str, str]: + from routes.inference import get_llama_cpp_backend + from core.inference import get_inference_backend + + llama = get_llama_cpp_backend() + if llama.is_loaded: + model = str(getattr(llama, "model_identifier", "") or "").strip() + variant = str(getattr(llama, "hf_variant", "") or "").strip() + return True, model, variant + + backend = get_inference_backend() + active_model = str(getattr(backend, "active_model_name", "") or "").strip() + if active_model: + return True, active_model, "" + return False, "", "" + + +def _ensure_selected_local_model_loaded( + recipe: dict[str, Any], local_provider_names: set[str] +) -> None: + model_loaded, active_model, active_variant = _loaded_local_model_identity() + if not model_loaded: + raise ValueError( + "No model loaded in Chat. Load a model first, then run the recipe." + ) + + selection = _single_used_local_model_selection(recipe, local_provider_names) + if selection is None: + return + + target, gguf_variant = selection + variant_matches = not gguf_variant or active_variant == gguf_variant + if active_model.lower() != target.lower() or not variant_matches: + selected = f"{target} ({gguf_variant})" if gguf_variant else target + active = ( + f"{active_model} ({active_variant})" if active_variant else active_model + ) + raise ValueError( + "Selected local model is not loaded. " + f"Selected {selected}; active {active or 'none'}. " + "Load the selected model again, then run the recipe." + ) + + def _inject_local_structured_response_format( recipe: dict[str, Any], local_provider_names: set[str] ) -> None: @@ -238,24 +321,12 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona token = "" internal_key_id: Optional[int] = None if local_names & referenced_providers: - # Verify a model is loaded. - # NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded - # or swapped after this check but before the recipe subprocess calls /v1. - # The inference endpoint returns a clear 400 in that case. - # - # Imports are deferred to avoid circular dependencies with inference modules. - from routes.inference import get_llama_cpp_backend - from core.inference import get_inference_backend - - llama = get_llama_cpp_backend() - model_loaded = llama.is_loaded - if not model_loaded: - backend = get_inference_backend() - model_loaded = bool(backend.active_model_name) - if not model_loaded: - raise ValueError( - "No model loaded in Chat. Load a model first, then run the recipe." - ) + # Verify the selected local model is loaded before minting a workflow + # key. This still remains a point-in-time singleton-backend check + # (TOCTOU): a future generation token should bind frontend load and + # job creation, and the inference endpoint returns a clear 400 if the + # model is later unloaded or swapped before the subprocess calls /v1. + _ensure_selected_local_model_loaded(recipe, local_names) from auth import storage # deferred: avoids circular import @@ -287,12 +358,12 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona providers[i].pop("extra_body", None) # Force skip_health_check on any model_config that references a local - # provider. The local /v1/models endpoint only lists the real loaded - # model (e.g. "unsloth/llama-3.2-1b") and not the placeholder "local" - # that the recipe sends as the model id, so data_designer's pre-flight - # health check would otherwise fail before the first completion call. - # The backend route ignores the model id field in chat completions, so - # skipping the check is safe. + # provider. The frontend now sends the explicit selected local model id, + # but llama-server's /v1/models response can still differ from that id + # for local paths, cache aliases, and GGUF variant loads. The recipe run + # has already gated on a loaded local inference backend above, so the + # data_designer model-list health check would be redundant and can reject + # valid local selections. for mc in recipe.get("model_configs", []): if not isinstance(mc, dict): continue @@ -319,7 +390,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona tpl_kwargs = extra_body.get("chat_template_kwargs") if not isinstance(tpl_kwargs, dict): tpl_kwargs = {} - tpl_kwargs.setdefault("enable_thinking", False) + tpl_kwargs["enable_thinking"] = False extra_body["chat_template_kwargs"] = tpl_kwargs params["extra_body"] = extra_body diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2b84b40d93..50fd283791 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -467,9 +467,19 @@ _TOOL_ACTION_NUDGE = ( " Do NOT output code blocks -- use the python tool instead." ) -# Regex for stripping leaked tool-call XML from assistant messages/stream +# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py +# split across the visible/DRAIN boundary. Four leak shapes: +# 1. well-formed `...` / `...` +# 2. orphan opening to EOF (close was DRAINED) +# 3. bare orphan close (open was DRAINED) +# 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".*?|.*?", + # 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, ) logger = get_logger(__name__) @@ -638,11 +648,16 @@ async def load_model( backend = get_inference_backend() llama_backend = get_llama_cpp_backend() - if request.gguf_variant: + is_direct_gguf_request = model_identifier.lower().endswith(".gguf") + if request.gguf_variant or is_direct_gguf_request: + gguf_variant_matches = is_direct_gguf_request or bool( + llama_backend.hf_variant + and request.gguf_variant + and llama_backend.hf_variant.lower() == request.gguf_variant.lower() + ) if ( llama_backend.is_loaded - and llama_backend.hf_variant - and llama_backend.hf_variant.lower() == request.gguf_variant.lower() + and gguf_variant_matches and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() # Match runtime settings too so Apply isn't dropped (#5401). @@ -651,7 +666,8 @@ async def load_model( and getattr(llama_backend, "_audio_probed", True) ): logger.info( - f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload" + "Model already loaded (GGUF): " + f"{model_log_label} variant={request.gguf_variant or llama_backend.hf_variant}, skipping reload" ) inference_config = load_inference_config(llama_backend.model_identifier) @@ -1405,6 +1421,7 @@ async def get_status( _audio_type = getattr(llama_backend, "_audio_type", None) return InferenceStatusResponse( active_model = _display_model_id, + model_identifier = None if _native_grant_backed else _model_id, is_vision = llama_backend.is_vision, is_gguf = True, gguf_variant = llama_backend.hf_variant, @@ -1467,6 +1484,7 @@ async def get_status( return InferenceStatusResponse( active_model = backend.active_model_name, + model_identifier = backend.active_model_name, is_vision = is_vision, is_gguf = False, is_audio = is_audio, @@ -1729,6 +1747,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. @@ -1741,6 +1760,12 @@ def _build_external_messages( see ``_INPUT_DOCUMENT_PROVIDERS``). For every other provider the part is stripped so the unknown content type doesn't reach generic /chat/completions passthrough and 400 the request. + - `reasoning`: OpenAI-only Responses reasoning item paired with a + prior tool output. Forwarded ONLY when provider_type=="openai" + so follow-up image edits can replay the required reasoning item. + - `image_generation_call`: OpenAI-only Responses image reference. + Forwarded ONLY when provider_type=="openai" so follow-up image + edits can reference prior generated images. - `compaction`: Anthropic-only synthetic part (round-trips server-side compaction state). Forwarded ONLY when provider_type=="anthropic"; stripped for every other provider so the unknown part doesn't @@ -1749,14 +1774,172 @@ 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: @@ -1769,6 +1952,30 @@ def _build_external_messages( "image_url": {"url": part.image_url.url}, } ) + elif ( + part.type == "reasoning" and openai and msg.role == "assistant" + ): + reasoning: dict[str, Any] = { + "type": "reasoning", + "id": part.id, + "summary": part.summary, + } + if part.status: + reasoning["status"] = part.status + parts.append(reasoning) + elif ( + part.type == "image_generation_call" + and openai + and msg.role == "assistant" + ): + # ExternalProviderClient maps this onto a top-level + # Responses input item after the current user prompt, + # or onto `previous_response_id` when response_id is + # available from the prior Responses turn. + image_ref = {"type": "image_generation_call", "id": part.id} + if getattr(part, "response_id", None): + image_ref["response_id"] = part.response_id + parts.append(image_ref) elif part.type == "input_document" and document_provider: # ExternalProviderClient maps this onto # Anthropic's `document` or OpenAI Responses' @@ -1790,7 +1997,27 @@ def _build_external_messages( # provider would 400 on the unknown part, so # gate by provider_type. parts.append({"type": "compaction", "content": part.content}) - result.append({"role": msg.role, "content": 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 + 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 -- @@ -1801,14 +2028,57 @@ def _build_external_messages( for p in msg.content: if p.type == "text": preserved.append({"type": "text", "text": p.text}) + elif p.type == "reasoning" and openai and msg.role == "assistant": + reasoning: dict[str, Any] = { + "type": "reasoning", + "id": p.id, + "summary": p.summary, + } + if p.status: + reasoning["status"] = p.status + preserved.append(reasoning) + elif ( + p.type == "image_generation_call" + and openai + and msg.role == "assistant" + ): + image_ref = {"type": "image_generation_call", "id": p.id} + if getattr(p, "response_id", None): + image_ref["response_id"] = p.response_id + preserved.append(image_ref) elif p.type == "compaction" and anthropic: preserved.append({"type": "compaction", "content": p.content}) + if msg.role == "assistant" and not preserved: + continue 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 @@ -1883,6 +2153,7 @@ async def _proxy_to_external_provider( payload.messages, _supports_vision, provider_type = provider_type, + base_url = base_url, ) client = ExternalProviderClient( @@ -1891,6 +2162,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, @@ -1899,7 +2178,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, @@ -1908,6 +2187,9 @@ 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, ) try: @@ -2408,17 +2690,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: # Preserve client-supplied order so prioritised tools # (e.g. search_knowledge_base when RAG is on) appear first. _by_name = {t["function"]["name"]: t for t in ALL_TOOLS} @@ -2431,6 +2725,19 @@ async def openai_chat_completions( tools_to_use, payload.rag_scope ) + 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 @@ -2909,8 +3216,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 @@ -2918,9 +3232,11 @@ 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: _by_name = {t["function"]["name"]: t for t in ALL_TOOLS} _sf_tools_to_use = [ _by_name[name] for name in payload.enabled_tools if name in _by_name @@ -2931,6 +3247,16 @@ async def openai_chat_completions( _sf_tools_to_use, payload.rag_scope ) + 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 @@ -4464,7 +4790,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. @@ -5255,6 +5591,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. @@ -5271,8 +5711,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: @@ -5321,8 +5763,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..a7501d1691 --- /dev/null +++ b/studio/backend/routes/mcp_servers.py @@ -0,0 +1,223 @@ +# 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, + list_tools_async, + parse_server_headers, +) +from models.mcp_servers import ( + McpServerCreate, + McpServerProbeResult, + McpServerResponse, + McpServerTestRequest, + McpServerUpdate, +) +from storage import mcp_servers_db + +logger = structlog.get_logger(__name__) + +router = APIRouter() + + +_PROBE_TIMEOUT_SECONDS = 8.0 +# When OAuth probes need to open a browser, wait long enough for the user to +# sign in. Matches fastmcp's default OAuth callback_timeout (300 s) + slack. +_OAUTH_PROBE_TIMEOUT_SECONDS = 305.0 + + +def _validate_url(url: str) -> str: + trimmed = (url or "").strip() + if not trimmed: + raise HTTPException(status_code = 400, detail = "url must not be empty") + 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) + + 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 = payload.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 + 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") + # 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") + + use_oauth = bool(server.get("use_oauth")) + try: + tools = await list_tools_async( + url = server["url"], + headers = parse_server_headers(server), + timeout = _OAUTH_PROBE_TIMEOUT_SECONDS + if use_oauth + else _PROBE_TIMEOUT_SECONDS, + 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 = _OAUTH_PROBE_TIMEOUT_SECONDS + if payload.use_oauth + else _PROBE_TIMEOUT_SECONDS, + 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 d5ccc49022..5ddf404370 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -9,6 +9,7 @@ Works independently and can be moved to any directory. import os import sys from pathlib import Path +from typing import Optional # Suppress annoying C-level dependency warnings globally (e.g. SwigPyPacked) os.environ["PYTHONWARNINGS"] = "ignore" @@ -18,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 @@ -512,10 +523,94 @@ _server = None _shutdown_event = None +_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist" + + +def _iter_frontend_fallback_candidates() -> "list[Path]": + """Yield `studio/frontend/dist` paths to try when the default is missing. + + Covers PATH-shadowed binaries whose __file__ resolves into a + site-packages tree that never received a vite build (e.g. plain + `pip install unsloth` from PyPI). + """ + import ast + import re + + out: list[Path] = [] + home_str = ( + os.environ.get("UNSLOTH_STUDIO_HOME") + or os.environ.get("STUDIO_HOME") + or str(Path.home() / ".unsloth" / "studio") + ) + venv_dir = Path(home_str).expanduser() / "unsloth_studio" + # Installer venv site-packages. + for pattern in ( + "lib/python*/site-packages/studio/frontend/dist", + "Lib/site-packages/studio/frontend/dist", + ): + out.extend(venv_dir.glob(pattern)) + # Editable source roots referenced from the installer venv. + for sp_pattern in ("lib/python*/site-packages", "Lib/site-packages"): + for sp in venv_dir.glob(sp_pattern): + for finder in sp.glob("__editable___*_finder.py"): + try: + src = finder.read_text(encoding = "utf-8") + except OSError: + continue + # Tolerate single- or multi-line dict literals; [^}]* still + # rejects nested dicts, which the setuptools template never + # emits for editable installs. + m = re.search( + r"^MAPPING\s*(?::[^=]*)?=\s*(\{[^}]*\})", src, re.M | re.S + ) + if not m: + continue + try: + mapping = ast.literal_eval(m.group(1)) + except (SyntaxError, ValueError): + continue + # Defensive: literal_eval can return a set / list / None if the + # matched literal is not a dict (regex captures `{...}`). + if not isinstance(mapping, dict): + continue + studio_pkg = mapping.get("studio") + if studio_pkg: + out.append(Path(studio_pkg) / "frontend" / "dist") + return out + + +def _resolve_frontend_path(frontend_path: Path) -> tuple[Optional[Path], list[Path]]: + """Pick a frontend dir that actually contains `index.html`. + + Returns (chosen, attempted). `chosen` is None if nothing servable was + found; `attempted` is the full ordered list for diagnostics. + """ + attempted: list[Path] = [] + seen: set[Path] = set() + + def _try(p: Path) -> bool: + try: + key = p.resolve() + except OSError: + key = p + if key in seen: + return False + seen.add(key) + attempted.append(p) + return (p / "index.html").is_file() + + if _try(Path(frontend_path)): + return attempted[-1], attempted + for alt in _iter_frontend_fallback_candidates(): + if _try(alt): + return attempted[-1], attempted + return None, attempted + + def run_server( host: str = "127.0.0.1", port: int = 8888, - frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist", + frontend_path: Path = _DEFAULT_FRONTEND_PATH, silent: bool = False, api_only: bool = False, llama_parallel_slots: int = 1, @@ -584,14 +679,48 @@ def run_server( print("=" * 50) print("") - # Setup frontend if path provided (skip in api-only mode) + # Setup frontend if path provided (skip in api-only mode). + # Falls back through alternate locations if the default lacks a built + # dist; errors out loudly rather than silently serving 404 on `/`. if frontend_path and not api_only: - if setup_frontend(app, frontend_path): + chosen, attempted = _resolve_frontend_path(Path(frontend_path)) + if chosen is not None and setup_frontend(app, chosen): if not silent: - print(f"[OK] Frontend loaded from {frontend_path}") + # Resolve so logs always show an absolute path for support. + try: + display = chosen.resolve() + except OSError: + display = chosen + print(f"[OK] Frontend loaded from {display}") else: - if not silent: - print(f"[WARNING] Frontend not found at {frontend_path}") + home_str = ( + os.environ.get("UNSLOTH_STUDIO_HOME") + or os.environ.get("STUDIO_HOME") + or str(Path.home() / ".unsloth" / "studio") + ) + # Windows ships the user-facing shim at $STUDIO_HOME/bin/unsloth.exe + # (a hardlink to the venv exe); Linux/macOS use the venv binary + # at $STUDIO_HOME/unsloth_studio/bin/unsloth. + home = Path(home_str).expanduser() + if sys.platform == "win32": + installer_bin = home / "bin" / "unsloth.exe" + else: + installer_bin = home / "unsloth_studio" / "bin" / "unsloth" + tried_lines = "\n".join(f" - {p}" for p in attempted) or " (none)" + raise SystemExit( + "[ERROR] Studio frontend build not found.\n" + f"Tried:\n{tried_lines}\n" + "\n" + "Likely cause: another 'unsloth' on PATH is shadowing the " + "installer's binary and points at a site-packages tree with " + "no built dist.\n" + "\n" + "Fix one of:\n" + f" - run the installer's binary directly: {installer_bin} studio\n" + " - pass --frontend \n" + " - pass --api-only to skip serving the web UI\n" + " - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh" + ) # Resolve once; shared by the log rewrite and the banner. display_host = _resolve_external_ip() if host == "0.0.0.0" else host @@ -718,7 +847,7 @@ if __name__ == "__main__": parser.add_argument( "--frontend", type = str, - default = Path(__file__).resolve().parent.parent / "frontend" / "dist", + default = _DEFAULT_FRONTEND_PATH, help = "Path to frontend build", ) parser.add_argument("--silent", action = "store_true", help = "Suppress output") @@ -727,11 +856,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_anthropic_citations.py b/studio/backend/tests/test_anthropic_citations.py new file mode 100644 index 0000000000..ab5ba10b56 --- /dev/null +++ b/studio/backend/tests/test_anthropic_citations.py @@ -0,0 +1,353 @@ +# 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 Anthropic ``citations_delta`` handling in the streaming proxy. + +Verifies the proxy injects inline ``[N]`` markers after cited text, +dedupes by type-specific anchor (char_location, page_location, +content_block_location, search_result_location), forwards a synthetic +``document_citations`` tool_event at message_stop, and stays inert when +no citations_delta events fire. See +https://platform.claude.com/docs/en/build-with-claude/citations +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _sse(events: list[dict]) -> bytes: + out = [] + for e in events: + ev = e.get("type", "message") + out.append(f"event: {ev}\ndata: {json.dumps(e)}\n\n") + return "".join(out).encode("utf-8") + + +def _capture(monkeypatch, events: list[dict]) -> list[str]: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + content = _sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + lines: list[str] = [] + + async def run(): + client = _make_client() + try: + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "what color is grass?"}], + model = "claude-opus-4-7", + max_tokens = 64, + ): + lines.append(line) + finally: + await client.close() + + _drive(run()) + return lines + + +def _message_start() -> dict: + return { + "type": "message_start", + "message": { + "id": "m1", + "content": [], + "model": "claude-opus-4-7", + "role": "assistant", + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + + +def _content_block_start_text() -> dict: + return { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + } + + +def _text_delta(text: str, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "text_delta", "text": text}, + } + + +def _citations_delta(citation: dict, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "citations_delta", "citation": citation}, + } + + +def _content_block_stop(index: int = 0) -> dict: + return {"type": "content_block_stop", "index": index} + + +def _message_delta_end() -> dict: + return {"type": "message_delta", "delta": {"stop_reason": "end_turn"}} + + +def _message_stop() -> dict: + return {"type": "message_stop"} + + +def _joined(lines: list[str]) -> str: + return "\n".join(lines) + + +def test_no_citations_stream_unchanged(monkeypatch): + """Plain text streams pass through with no inline markers and no + document_citations tool_event.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Grass is green." in body + assert "document_citations" not in body + assert "[1]" not in body + + +def test_single_char_location_emits_inline_marker(monkeypatch): + cit = { + "type": "char_location", + "cited_text": "The grass is green.", + "document_index": 0, + "document_title": "Example", + "start_char_index": 0, + "end_char_index": 20, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Grass is green." in body + assert "[1]" in body, body + assert "document_citations" in body, body + assert '"document_index": 0' in body, body + assert "_key" not in body, body + + +def test_duplicate_citation_dedupes_to_same_number(monkeypatch): + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Example", + "start_char_index": 0, + "end_char_index": 20, + "cited_text": "The grass is green.", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass."), + _citations_delta(cit), + _text_delta(" Still green."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert body.count("[1]") == 2, body + citation_blob = body[body.index("document_citations") :] + assert citation_blob.count('"start_char_index"') == 1, citation_blob + + +def test_distinct_sources_get_distinct_numbers(monkeypatch): + cit1 = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc A", + "start_char_index": 0, + "end_char_index": 5, + } + cit2 = { + "type": "page_location", + "document_index": 1, + "document_title": "Doc B", + "start_page_number": 3, + "end_page_number": 4, + } + cit3 = { + "type": "content_block_location", + "document_index": 2, + "document_title": "Doc C", + "start_block_index": 0, + "end_block_index": 1, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("First"), + _citations_delta(cit1), + _text_delta(" Second"), + _citations_delta(cit2), + _text_delta(" Third"), + _citations_delta(cit3), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body and "[2]" in body and "[3]" in body, body + assert body.index("[1]") < body.index("[2]") < body.index("[3]") + + +def test_search_result_location_supported(monkeypatch): + cit = { + "type": "search_result_location", + "document_index": 0, + "document_title": "Anthropic Search Results", + "source": "https://example.com/doc.html", + "start_block_index": 0, + "end_block_index": 1, + "cited_text": "blah", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Some sourced fact."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body + assert "search_result_location" in body + assert "example.com/doc.html" in body + + +def test_same_start_different_end_offsets_get_distinct_numbers(monkeypatch): + """Same start_char_index + different end_char_index = distinct spans, + so they must get distinct footnote numbers (ranges use exclusive end).""" + cit_a = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 100, + "end_char_index": 150, + "cited_text": "first half", + } + cit_b = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 100, + "end_char_index": 250, + "cited_text": "wider span", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("A "), + _citations_delta(cit_a), + _text_delta(" and B "), + _citations_delta(cit_b), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + assert "[2]" in body, body + + +def test_search_result_location_different_indices_get_distinct_numbers(monkeypatch): + """Same source + different search_result_index = distinct footnotes + (matches the Anthropic search-result citation contract).""" + cit_a = { + "type": "search_result_location", + "search_result_index": 0, + "source": "https://example.com/result.html", + "title": "Result", + "start_block_index": 0, + "end_block_index": 1, + "cited_text": "first", + } + cit_b = { + "type": "search_result_location", + "search_result_index": 1, + "source": "https://example.com/result.html", + "title": "Result", + "start_block_index": 0, + "end_block_index": 1, + "cited_text": "second", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("A "), + _citations_delta(cit_a), + _text_delta(" and B "), + _citations_delta(cit_b), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + assert "[2]" in body, body diff --git a/studio/backend/tests/test_anthropic_citations_edge.py b/studio/backend/tests/test_anthropic_citations_edge.py new file mode 100644 index 0000000000..be1b5f7922 --- /dev/null +++ b/studio/backend/tests/test_anthropic_citations_edge.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case tests for Anthropic ``citations_delta`` handling. + +Complements ``test_anthropic_citations.py``. Covers malformed payloads, +unusual orderings, mixed citation types, and the ``citations: +{enabled: true}`` opt-in attached to translated ``input_document`` +blocks. See +https://platform.claude.com/docs/en/build-with-claude/citations and +https://platform.claude.com/docs/en/build-with-claude/search-results. +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +# ── shared SSE harness ─────────────────────────────────────── + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _sse(events: list[dict]) -> bytes: + out = [] + for e in events: + ev = e.get("type", "message") + out.append(f"event: {ev}\ndata: {json.dumps(e)}\n\n") + return "".join(out).encode("utf-8") + + +def _capture( + monkeypatch, + events: list[dict], + *, + messages: list[dict] | None = None, + captured_body: dict | None = None, +) -> list[str]: + """Drive ``stream_chat_completion`` against a mocked Anthropic + response and return the SSE lines. Pass ``captured_body`` to also + capture the outgoing request body for assertions on the translated + Anthropic shape. + """ + + def handler(request: httpx.Request) -> httpx.Response: + if captured_body is not None: + try: + captured_body.update(json.loads(request.content.decode("utf-8"))) + except Exception: # pragma: no cover -- diagnostic only + pass + return httpx.Response( + 200, + content = _sse(events), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + lines: list[str] = [] + + async def run(): + client = _make_client() + try: + async for line in client.stream_chat_completion( + messages = messages + or [{"role": "user", "content": "what color is grass?"}], + model = "claude-opus-4-7", + max_tokens = 64, + ): + lines.append(line) + finally: + await client.close() + + _drive(run()) + return lines + + +def _message_start() -> dict: + return { + "type": "message_start", + "message": { + "id": "m1", + "content": [], + "model": "claude-opus-4-7", + "role": "assistant", + "stop_reason": None, + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + + +def _content_block_start_text() -> dict: + return { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + } + + +def _text_delta(text: str, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "text_delta", "text": text}, + } + + +def _citations_delta(citation: dict, index: int = 0) -> dict: + return { + "type": "content_block_delta", + "index": index, + "delta": {"type": "citations_delta", "citation": citation}, + } + + +def _content_block_stop(index: int = 0) -> dict: + return {"type": "content_block_stop", "index": index} + + +def _message_delta_end() -> dict: + return {"type": "message_delta", "delta": {"stop_reason": "end_turn"}} + + +def _message_stop() -> dict: + return {"type": "message_stop"} + + +def _joined(lines: list[str]) -> str: + return "\n".join(lines) + + +def _citation_payload(body: str) -> dict: + """Pull the ``document_citations`` synthetic tool_event from the + SSE body and return its payload. Raises if absent.""" + assert "document_citations" in body, body + for line in body.splitlines(): + if not line.startswith("data: "): + continue + try: + payload = json.loads(line[len("data: ") :]) + except json.JSONDecodeError: + continue + tool_event = payload.get("_toolEvent") if isinstance(payload, dict) else None + if ( + isinstance(tool_event, dict) + and tool_event.get("type") == "document_citations" + ): + return tool_event + raise AssertionError("document_citations event not parsed out of SSE body") + + +# ── edge cases ─────────────────────────────────────────────── + + +def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch): + """citations_delta before any text_delta must not crash; marker + lands at the start of the block.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "X", + "start_char_index": 0, + "end_char_index": 5, + "cited_text": "x", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _citations_delta(cit), + _text_delta("hello"), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + assert "document_citations" in body, body + + +def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch): + """Non-dict ``delta.citation`` must not crash, emit a marker, or + poison the document_citations list.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Hello."), + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "citations_delta", "citation": "not-a-dict"}, + }, + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Hello." in body + assert "[1]" not in body + assert "document_citations" not in body + + +def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch): + """Missing ``citation`` field is treated like a non-dict citation: + skip without crashing.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Hello."), + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "citations_delta"}, + }, + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "Hello." in body + assert "[1]" not in body + assert "document_citations" not in body + + +def test_char_location_with_reversed_indices_does_not_crash(monkeypatch): + """Malformed char_location with reversed indices must not crash; + the dedup key accepts any int pair and still surfaces a footnote.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 300, + "end_char_index": 50, + "cited_text": "?", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Weird."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + payload = _citation_payload(body) + assert payload["citations"][0]["start_char_index"] == 300 + assert payload["citations"][0]["end_char_index"] == 50 + + +def test_page_location_missing_document_index_does_not_crash(monkeypatch): + """page_location missing ``document_index`` still produces a + footnote; dedup key falls back to ``None`` for the missing field.""" + cit = { + "type": "page_location", + "document_title": "Untitled PDF", + "start_page_number": 1, + "end_page_number": 2, + "cited_text": "p1", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("From the PDF:"), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + payload = _citation_payload(body) + assert payload["citations"][0].get("document_index") is None + + +def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypatch): + """content_block_location with string block indices must not crash; + dedup key tolerates non-int values.""" + cit = { + "type": "content_block_location", + "document_index": 0, + "document_title": "Custom", + "start_block_index": "0", + "end_block_index": "1", + "cited_text": "anything", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Cite."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body, body + payload = _citation_payload(body) + assert payload["citations"][0]["start_block_index"] == "0" + + +def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch): + """Unknown citation ``type`` (forward-compat) still dedupes: + identical ones collapse, differing ones get distinct numbers.""" + cit_a = { + "type": "future_shape_location", + "anchor": "abc", + "cited_text": "blah", + } + cit_b = { + "type": "future_shape_location", + "anchor": "xyz", + "cited_text": "blah", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("A"), + _citations_delta(cit_a), + _text_delta(" again"), + _citations_delta(cit_a), + _text_delta(" B"), + _citations_delta(cit_b), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + # cit_a dedupes onto [1], cit_b gets [2]. + assert body.count("[1]") == 2, body + assert body.count("[2]") == 1, body + payload = _citation_payload(body) + assert len(payload["citations"]) == 2 + + +def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch): + """char_location and page_location on the same document_index are + distinct shapes; dedup key uses citation type as its first slot.""" + cit_char = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 10, + } + cit_page = { + "type": "page_location", + "document_index": 0, + "document_title": "Doc", + "start_page_number": 1, + "end_page_number": 2, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("char-cite"), + _citations_delta(cit_char), + _text_delta(" page-cite"), + _citations_delta(cit_page), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body and "[2]" in body, body + payload = _citation_payload(body) + assert len(payload["citations"]) == 2 + + +def test_cited_text_is_preserved_in_synthetic_event(monkeypatch): + """``cited_text`` must survive into the synthetic event so the + Sources panel can render it as a tooltip. Anthropic does not bill + cited_text against output tokens, so preserving it is free.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Trustworthy Doc", + "start_char_index": 0, + "end_char_index": 20, + "cited_text": "The grass is green.", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + payload = _citation_payload(body) + assert payload["citations"][0]["cited_text"] == "The grass is green." + + +def test_internal_key_field_never_leaks_to_client(monkeypatch): + """The internal ``_key`` dedup sentinel must be stripped before + the synthetic event is forwarded; it is not an Anthropic field.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 5, + "cited_text": "..", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("hi"), + _citations_delta(cit), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + payload = _citation_payload(body) + assert payload["citations"], payload + for c in payload["citations"]: + assert "_key" not in c, c + + +def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch): + """Footnote numbering is per-message, not per-content-block: + citations across separate blocks emit [1] then [2].""" + cit_a = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 5, + } + cit_b = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 100, + "end_char_index": 105, + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("first"), + _citations_delta(cit_a, index = 0), + _content_block_stop(0), + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "text", "text": ""}, + }, + _text_delta(" second", index = 1), + _citations_delta(cit_b, index = 1), + _content_block_stop(1), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "[1]" in body and "[2]" in body, body + assert body.index("[1]") < body.index("[2]") + payload = _citation_payload(body) + assert len(payload["citations"]) == 2 + + +def test_inline_marker_lands_after_text_run(monkeypatch): + """Inline ``[N]`` must land AFTER the cited text run: Anthropic + streams text then citation, so the proxy emits ``"...green.[1]"`` + not ``"[1]green"``.""" + cit = { + "type": "char_location", + "document_index": 0, + "document_title": "Doc", + "start_char_index": 0, + "end_char_index": 20, + "cited_text": "grass", + } + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Grass is green."), + _citations_delta(cit), + _text_delta(" Sky is blue."), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + grass = body.index("Grass is green.") + marker = body.index("[1]") + sky = body.index("Sky is blue.") + assert grass < marker < sky, body + + +def test_no_synthetic_event_when_only_text_deltas(monkeypatch): + """No citations_delta means no synthetic ``document_citations`` + event; Sources panel relies on absence to suppress the section.""" + lines = _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("Just some prose. "), + _text_delta("More prose."), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + ) + body = _joined(lines) + assert "document_citations" not in body + assert "[1]" not in body + + +def test_input_document_translation_enables_citations(monkeypatch): + """``input_document`` must translate to an Anthropic ``document`` + block carrying ``citations: {enabled: true}`` (both base64 and url + source branches) so upstream emits citations_delta.""" + captured_b64: dict = {} + _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("ok"), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + messages = [ + { + "role": "user", + "content": [ + { + "type": "input_document", + "file_data": "data:application/pdf;base64,QUJD", + "filename": "spec.pdf", + }, + {"type": "text", "text": "summarise"}, + ], + } + ], + captured_body = captured_b64, + ) + user_msg = captured_b64["messages"][0] + doc_block = next(p for p in user_msg["content"] if p.get("type") == "document") + assert doc_block["source"]["type"] == "base64", doc_block + assert doc_block.get("citations") == {"enabled": True}, doc_block + + captured_url: dict = {} + _capture( + monkeypatch, + [ + _message_start(), + _content_block_start_text(), + _text_delta("ok"), + _content_block_stop(), + _message_delta_end(), + _message_stop(), + ], + messages = [ + { + "role": "user", + "content": [ + { + "type": "input_document", + "file_url": "https://example.com/doc.pdf", + "filename": "doc.pdf", + }, + {"type": "text", "text": "summarise"}, + ], + } + ], + captured_body = captured_url, + ) + user_msg = captured_url["messages"][0] + doc_block = next(p for p in user_msg["content"] if p.get("type") == "document") + assert doc_block["source"]["type"] == "url", doc_block + assert doc_block.get("citations") == {"enabled": True}, doc_block + + +# ── cited_text truncation + safe-url citation conversion ──────── + + +def test_cited_text_truncated_in_synthetic_event(monkeypatch): + """``cited_text`` is capped server-side so multi-KB spans do not + balloon the SSE payload.""" + from core.inference.external_provider import _CITED_TEXT_MAX_LEN + + long_quote = "x" * (_CITED_TEXT_MAX_LEN + 4000) + events = [ + { + "type": "message_start", + "message": { + "id": "msg_1", + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "claim "}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "citations_delta", + "citation": { + "type": "char_location", + "document_index": 0, + "document_title": "doc", + "start_char_index": 0, + "end_char_index": 5, + "cited_text": long_quote, + }, + }, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 1}, + }, + {"type": "message_stop"}, + ] + chunks = _capture(monkeypatch, events) + tool_events = [c for c in chunks if "_toolEvent" in c and "document_citations" in c] + assert tool_events, "no document_citations tool event" + payload = json.loads(tool_events[0].split("data: ", 1)[1]) + cited = payload["_toolEvent"]["citations"][0]["cited_text"] + assert len(cited) <= _CITED_TEXT_MAX_LEN + 1, len(cited) + assert cited.endswith("…") 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_fast_mode_and_refusal.py b/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py new file mode 100644 index 0000000000..e7e5ec64d4 --- /dev/null +++ b/studio/backend/tests/test_anthropic_fast_mode_and_refusal.py @@ -0,0 +1,164 @@ +# 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 Anthropic fast-mode wiring and streaming refusal handling. + +fast_mode=True on Opus 4.6/4.7 attaches the ``fast-mode-2026-02-01`` +beta header and sets ``speed: "fast"``; unsupported models drop both. +Streaming ``stop_reason: "refusal"`` surfaces a user notice before the +``content_filter`` finish chunk. +https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _empty_message_sse() -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"claude-opus-4-7","role":"assistant",' + b'"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"end_turn"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _refusal_sse() -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"claude-opus-4-7","role":"assistant",' + b'"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"Hello."}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"refusal"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]: + """Install a MockTransport, drive one streamed call, return body+lines.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = sse or _empty_message_sse(), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + out_lines: list[str] = [] + + async def run(): + client = _make_client() + try: + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = kwargs.get("model", "claude-opus-4-7"), + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + fast_mode = kwargs.get("fast_mode"), + ): + out_lines.append(line) + finally: + await client.close() + + _drive(run()) + return captured, out_lines + + +def test_fast_mode_attaches_beta_header_and_speed_on_opus_4_7(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7") + assert cap["body"].get("speed") == "fast", cap["body"] + beta = cap["headers"].get("anthropic-beta", "") + assert "fast-mode-2026-02-01" in beta, beta + + +def test_fast_mode_attaches_beta_header_and_speed_on_opus_4_6(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-6") + assert cap["body"].get("speed") == "fast", cap["body"] + assert "fast-mode-2026-02-01" in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_dropped_on_sonnet(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-sonnet-4-6") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_dropped_on_haiku(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-haiku-4-5") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_dropped_on_older_opus(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-5") + assert "speed" not in cap["body"], cap["body"] + + +def test_fast_mode_false_does_not_attach_header_or_field(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = False) + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_none_does_not_attach_header_or_field(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = None) + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_refusal_emits_user_facing_notice_and_content_filter_finish(monkeypatch): + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + # User-visible refusal notice. + assert "stopped by Anthropic's safety classifier" in body, body + # OpenAI-spec finish_reason mapping. + assert '"finish_reason": "content_filter"' in body, body + # Original deltas preserved before the refusal supplement. + assert "Hello." in body, body + + +def test_refusal_emits_tool_event_for_chat_adapter_drop(monkeypatch): + """Refused turns emit an out-of-band `_toolEvent` that the chat-adapter + latches into assistant `metadata.custom.anthropicRefusal`, driving + the next-request prune. Tool event (not text) prevents spoofing. + """ + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + assert '"_toolEvent": {"type": "anthropic_refusal"}' in body, body + # Visible refusal text must not embed a sentinel that could spoof + # a context reset if echoed by another assistant message. + assert "studio:anthropic-refusal" not in body, body diff --git a/studio/backend/tests/test_anthropic_fast_mode_edge.py b/studio/backend/tests/test_anthropic_fast_mode_edge.py new file mode 100644 index 0000000000..0052cb94ad --- /dev/null +++ b/studio/backend/tests/test_anthropic_fast_mode_edge.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case coverage for the Anthropic fast-mode + refusal wiring. + +Complements ``test_anthropic_fast_mode_and_refusal.py`` (happy path) +with dated snapshots, strict opt-in (future Opus families do not +auto-enable), multi-beta header merging, refusal stream ordering, and +the non-destruction guarantee for unset/None fast_mode. +""" + +import asyncio +import json +import re + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +def _make_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "anthropic", + base_url = "https://api.anthropic.com/v1", + api_key = "sk-ant-test", + ) + + +def _empty_message_sse(model: str = "claude-opus-4-7") -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"' + model.encode() + b'",' + b'"role":"assistant","stop_reason":null,"usage":' + b'{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"end_turn"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _refusal_sse(model: str = "claude-opus-4-7") -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"' + model.encode() + b'",' + b'"role":"assistant","stop_reason":null,"usage":' + b'{"input_tokens":1,"output_tokens":1}}}\n\n' + b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"Hello."}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"refusal"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def _capture(monkeypatch, sse: bytes = b"", **kwargs) -> tuple[dict, list[str]]: + """Install a MockTransport, drive one streamed call, return body+lines.""" + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = sse or _empty_message_sse(kwargs.get("model", "claude-opus-4-7")), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + out_lines: list[str] = [] + + async def run(): + client = _make_client() + try: + extra = {} + for key in ( + "enabled_tools", + "compaction_threshold", + "fast_mode", + ): + if key in kwargs: + extra[key] = kwargs[key] + async for line in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = kwargs.get("model", "claude-opus-4-7"), + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + **extra, + ): + out_lines.append(line) + finally: + await client.close() + + _drive(run()) + return captured, out_lines + + +# ──────────────────────────── dated snapshot prefix ──────────────────────────── +def test_fast_mode_attaches_on_dated_opus_4_7_snapshot(monkeypatch): + """Dated snapshot ``claude-opus-4-7-2026-02-01`` must match the prefix.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7-2026-02-01") + assert cap["body"].get("speed") == "fast", cap["body"] + assert "fast-mode-2026-02-01" in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_attaches_on_dated_opus_4_6_snapshot(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-6-2026-02-01") + assert cap["body"].get("speed") == "fast", cap["body"] + assert "fast-mode-2026-02-01" in cap["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── strict opt-in semantics ──────────────────────────── +def test_fast_mode_does_not_auto_enable_on_future_opus_4_8(monkeypatch): + """Future ``claude-opus-4-8`` must not auto-enable; opt-in per family.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-8") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_does_not_auto_enable_on_future_opus_5(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-5") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_does_not_auto_enable_on_sonnet_dated_snapshot(monkeypatch): + """Sonnet snapshots share the compaction prefix but not fast_mode.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-sonnet-4-6-2026-02-01") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── beta header merge ──────────────────────────── +def _beta_parts(headers: dict) -> list[str]: + raw = headers.get("anthropic-beta", "") + return [p.strip() for p in raw.split(",") if p.strip()] + + +def test_fast_mode_merges_with_code_execution_beta(monkeypatch): + """fast_mode + code_execution -> two comma-separated betas, no overwrite.""" + cap, _ = _capture( + monkeypatch, + fast_mode = True, + model = "claude-opus-4-7", + enabled_tools = ["code_execution"], + ) + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts, cap["headers"] + assert any(p.startswith("code-execution-") for p in parts), cap["headers"] + # No duplicates. + assert len(parts) == len(set(parts)), parts + + +def test_fast_mode_merges_with_compaction_beta(monkeypatch): + """fast_mode + compaction_threshold >= 50K -> both betas present.""" + cap, _ = _capture( + monkeypatch, + fast_mode = True, + model = "claude-opus-4-7", + compaction_threshold = 100_000, + ) + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts, cap["headers"] + assert "compact-2026-01-12" in parts, cap["headers"] + + +def test_fast_mode_merges_with_code_execution_and_compaction(monkeypatch): + """Three betas coexist in one comma-separated header, no duplicates.""" + cap, _ = _capture( + monkeypatch, + fast_mode = True, + model = "claude-opus-4-7", + enabled_tools = ["code_execution"], + compaction_threshold = 100_000, + ) + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts + assert "compact-2026-01-12" in parts + assert any(p.startswith("code-execution-") for p in parts), parts + assert len(parts) >= 3 + assert len(parts) == len(set(parts)), parts + + +def test_fast_mode_beta_value_is_pinned(monkeypatch): + """Pin the exact beta tag ``fast-mode-2026-02-01`` from the docs.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7") + parts = _beta_parts(cap["headers"]) + assert "fast-mode-2026-02-01" in parts, parts + # Reject obvious typos. + assert not any(p.startswith("fastmode-") for p in parts), parts + assert not any("fast_mode" in p for p in parts), parts + + +# ──────────────────────────── non-destruction guarantee ──────────────────────────── +def test_fast_mode_unset_is_byte_identical_to_omitted(monkeypatch): + """``fast_mode=None`` must produce the same body/headers as omission.""" + cap_none, _ = _capture(monkeypatch, fast_mode = None, model = "claude-opus-4-7") + + # Re-run without passing fast_mode at all. + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode("utf-8")) + captured["headers"] = dict(request.headers) + return httpx.Response( + 200, + content = _empty_message_sse(), + headers = {"content-type": "text/event-stream"}, + ) + + monkeypatch.setattr( + ep_mod, + "_http_client", + httpx.AsyncClient(transport = httpx.MockTransport(handler)), + ) + + async def run(): + client = _make_client() + try: + async for _ in client.stream_chat_completion( + messages = [{"role": "user", "content": "hi"}], + model = "claude-opus-4-7", + temperature = 0.7, + top_p = 0.95, + max_tokens = 32, + ): + pass + finally: + await client.close() + + _drive(run()) + + assert cap_none["body"] == captured["body"], (cap_none["body"], captured["body"]) + # Headers can vary by httpx-injected fields (host, connection); compare + # the load-bearing ones. + for key in ("anthropic-version", "x-api-key", "content-type"): + assert cap_none["headers"].get(key) == captured["headers"].get(key), key + assert "anthropic-beta" not in cap_none["headers"] + assert "anthropic-beta" not in captured["headers"] + assert "speed" not in cap_none["body"] + assert "speed" not in captured["body"] + + +def test_fast_mode_false_on_opus_4_7_byte_identical_to_unset(monkeypatch): + """``fast_mode=False`` produces the same outbound shape as unset.""" + cap_false, _ = _capture(monkeypatch, fast_mode = False, model = "claude-opus-4-7") + assert "speed" not in cap_false["body"], cap_false["body"] + assert "fast-mode-2026-02-01" not in cap_false["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── refusal stream ordering ──────────────────────────── +def test_refusal_notice_appears_before_content_filter_chunk(monkeypatch): + """The notice content delta must precede the finish_reason chunk.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + notice_idx = next(i for i, l in enumerate(lines) if "stopped by Anthropic" in l) + filter_idx = next( + i for i, l in enumerate(lines) if '"finish_reason": "content_filter"' in l + ) + assert notice_idx < filter_idx, (notice_idx, filter_idx, lines) + + +def test_refusal_tool_event_emitted_exactly_once(monkeypatch): + """A single refusal emits the chat-adapter drop signal exactly once.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + count = body.count('"_toolEvent": {"type": "anthropic_refusal"}') + assert count == 1, (count, body) + + +def test_refusal_text_carries_no_html_sentinel(monkeypatch): + """Visible refusal text must not embed a ``studio:anthropic-refusal`` + sentinel; the drop signal rides _toolEvent only.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse()) + body = "\n".join(lines) + assert "studio:anthropic-refusal" not in body, body + + +def test_refusal_handling_works_on_sonnet_model(monkeypatch): + """Refusal handling is provider-side; Sonnet refusals must also surface.""" + _, lines = _capture( + monkeypatch, sse = _refusal_sse("claude-sonnet-4-6"), model = "claude-sonnet-4-6" + ) + body = "\n".join(lines) + assert "stopped by Anthropic's safety classifier" in body, body + assert '"_toolEvent": {"type": "anthropic_refusal"}' in body, body + assert '"finish_reason": "content_filter"' in body, body + + +def test_refusal_preserves_partial_assistant_text(monkeypatch): + """Partial deltas already streamed must precede the refusal notice.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + body = "\n".join(lines) + hello_idx = body.index("Hello.") + notice_idx = body.index("stopped by Anthropic") + assert hello_idx < notice_idx, (hello_idx, notice_idx) + + +def test_refusal_chunk_is_proper_openai_delta_shape(monkeypatch): + """The notice rides ``choices[0].delta.content`` (not a finish chunk); + OpenAI-spec clients treat it as ordinary streamed text.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + # Find the chunk that carries the refusal text. + notice_chunk = None + for line in lines: + if line.startswith("data: ") and "stopped by Anthropic" in line: + notice_chunk = json.loads(line[len("data: ") :]) + break + assert notice_chunk is not None, lines + choice = notice_chunk["choices"][0] + assert "delta" in choice and "content" in choice["delta"], notice_chunk + # Must NOT carry a finish_reason itself -- that comes on the next + # chunk. + assert choice.get("finish_reason") in (None,), notice_chunk + # Refusal text is plain-spoken; no embedded sentinel. + assert "studio:anthropic-refusal" not in choice["delta"]["content"] + + +def test_refusal_tool_event_chunk_shape(monkeypatch): + """Drop signal rides a Studio `_toolEvent` envelope (delta={}, + finish_reason=null); the frontend latches on + `_toolEvent.type == "anthropic_refusal"`.""" + _, lines = _capture(monkeypatch, sse = _refusal_sse(), model = "claude-opus-4-7") + refusal_chunk = None + for line in lines: + if line.startswith("data: ") and "anthropic_refusal" in line: + refusal_chunk = json.loads(line[len("data: ") :]) + break + assert refusal_chunk is not None, lines + assert refusal_chunk["_toolEvent"] == {"type": "anthropic_refusal"}, refusal_chunk + choice = refusal_chunk["choices"][0] + assert choice["delta"] == {}, refusal_chunk + assert choice["finish_reason"] is None, refusal_chunk + + +# ──────────────────────────── future-proofing ──────────────────────────── +def test_fast_mode_prefix_tuple_matches_capability_doc(monkeypatch): + """Tuple must exactly match the two families in the upstream docs: + https://platform.claude.com/docs/en/build-with-claude/fast-mode.""" + from core.inference.external_provider import _ANTHROPIC_FAST_MODE_PREFIXES + + assert set(_ANTHROPIC_FAST_MODE_PREFIXES) == { + "claude-opus-4-7", + "claude-opus-4-6", + }, _ANTHROPIC_FAST_MODE_PREFIXES + + +def test_fast_mode_speed_field_value_is_literal_fast(monkeypatch): + """Pin the wire value to the literal string ``"fast"``.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7") + assert cap["body"]["speed"] == "fast", cap["body"] + + +def test_fast_mode_dropped_on_opus_4_5_dated_snapshot(monkeypatch): + """Previous-family snapshots like ``claude-opus-4-5-2025-...`` must not match.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-5-2025-08-01") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_rejects_prefix_collision_4_70(monkeypatch): + """IDs like ``claude-opus-4-70`` / ``-4-7b`` must not match the prefix.""" + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-70") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_rejects_prefix_collision_4_7b(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-7b") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +def test_fast_mode_rejects_prefix_collision_4_6_extra(monkeypatch): + cap, _ = _capture(monkeypatch, fast_mode = True, model = "claude-opus-4-60") + assert "speed" not in cap["body"], cap["body"] + assert "fast-mode-2026-02-01" not in cap["headers"].get("anthropic-beta", "") + + +# ──────────────────────────── usage.speed propagation ──────────────────────────── +def _fast_speed_sse(model: str = "claude-opus-4-7", speed: str = "fast") -> bytes: + return ( + b'event: message_start\ndata: {"type":"message_start","message":' + b'{"id":"m1","content":[],"model":"' + model.encode() + b'",' + b'"role":"assistant","stop_reason":null,"usage":' + b'{"input_tokens":4,"output_tokens":1}}}\n\n' + b'event: content_block_start\ndata: {"type":"content_block_start",' + b'"index":0,"content_block":{"type":"text","text":""}}\n\n' + b'event: content_block_delta\ndata: {"type":"content_block_delta",' + b'"index":0,"delta":{"type":"text_delta","text":"hi"}}\n\n' + b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n' + b'event: message_delta\ndata: {"type":"message_delta",' + b'"delta":{"stop_reason":"end_turn"},' + b'"usage":{"output_tokens":5,"speed":"' + speed.encode() + b'"}}\n\n' + b'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + + +def test_usage_speed_propagates_to_final_usage_chunk_fast(monkeypatch): + """``usage.speed == "fast"`` from upstream must reach the Studio usage chunk.""" + _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "fast")) + usage_lines = [l for l in lines if l.startswith("data: ") and '"usage"' in l] + assert usage_lines, lines + parsed = [json.loads(l[len("data: ") :]) for l in usage_lines] + speeds = [p["usage"].get("speed") for p in parsed if "usage" in p] + assert "fast" in speeds, parsed + + +def test_usage_speed_propagates_to_final_usage_chunk_standard(monkeypatch): + _, lines = _capture(monkeypatch, sse = _fast_speed_sse(speed = "standard")) + parsed = [ + json.loads(l[len("data: ") :]) + for l in lines + if l.startswith("data: ") and '"usage"' in l + ] + speeds = [p["usage"].get("speed") for p in parsed if "usage" in p] + assert "standard" in speeds, parsed + + +def test_usage_speed_absent_when_anthropic_does_not_report(monkeypatch): + """Studio must not invent ``usage.speed`` when upstream omits it.""" + _, lines = _capture(monkeypatch) + parsed = [ + json.loads(l[len("data: ") :]) + for l in lines + if l.startswith("data: ") and '"usage"' in l + ] + for p in parsed: + usage = p.get("usage") or {} + assert "speed" not in usage, p diff --git a/studio/backend/tests/test_anthropic_web_fetch.py b/studio/backend/tests/test_anthropic_web_fetch.py index cdb5f6254c..88a922e7eb 100644 --- a/studio/backend/tests/test_anthropic_web_fetch.py +++ b/studio/backend/tests/test_anthropic_web_fetch.py @@ -2,26 +2,12 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Unit tests for Anthropic's server-side `web_fetch_20250910` tool -translation in `_stream_anthropic`. - -Covers: -- Request body: when ``enabled_tools=["web_fetch"]``, the outbound - ``tools`` array carries ``{"type":"web_fetch_20250910", - "name":"web_fetch", "max_uses":5}``. No beta header is required. -- Combined request: ``enabled_tools=["web_search","web_fetch", - "code_execution"]`` sends all three tool entries. -- Disabled by default: with ``enabled_tools=["web_search"]`` (or None), - the body does NOT carry a web_fetch entry. -- SSE translation (success): a `web_fetch` server_tool_use streaming - ``{"url": "..."}`` followed by a `web_fetch_tool_result` block with - a document source emits one ``tool_start`` and one ``tool_end`` - `_toolEvent`. The ``tool_start.arguments.url`` matches the fetched - URL and the ``tool_end.result`` carries the Title / URL / snippet - prefix the source-pill renderer expects. -- SSE translation (error): a `web_fetch_tool_error` with - ``error_code="url_not_accessible"`` renders as ``"Error: - url_not_accessible"`` in the tool_end result. +Unit tests for Anthropic's `web_fetch_20250910` / `web_fetch_20260209` +translation in ``_stream_anthropic``. Covers request body emission +(version picked by ``_anthropic_web_fetch_version``: ``_20260209`` for +Opus 4.6/4.7 + Sonnet 4.6, ``_20250910`` otherwise), combined tool +requests, off-by-default behavior, and SSE translation of success and +``url_not_accessible`` error paths into ``tool_start`` / ``tool_end``. """ import asyncio @@ -117,8 +103,9 @@ def test_web_fetch_tool_appended_to_request_body(monkeypatch): body = captured["body"] tools = body.get("tools") or [] + # claude-opus-4-7 routes web_fetch to _20260209 (dynamic filtering). assert { - "type": "web_fetch_20250910", + "type": "web_fetch_20260209", "name": "web_fetch", "max_uses": 5, } in tools @@ -157,13 +144,10 @@ def test_web_fetch_combined_with_web_search_and_code_execution(monkeypatch): tools = captured["body"].get("tools") or [] tool_types = [t.get("type") for t in tools] - # After PR 5679's per-model tool version dispatch landed, - # claude-opus-4-7 routes web_search to the _20260209 variant and - # code_execution to _20260120. web_fetch still hardcodes - # _20250910 today; see follow-up to thread it through - # _anthropic_web_fetch_version. + # claude-opus-4-7 routes web_search and web_fetch to _20260209 + # and code_execution to _20260120 (per PR 5679 dispatch). assert "web_search_20260209" in tool_types, tool_types - assert "web_fetch_20250910" in tool_types, tool_types + assert "web_fetch_20260209" in tool_types, tool_types assert "code_execution_20260120" in tool_types, tool_types # Code-execution still adds its beta flag; web_fetch must not # have accidentally stripped it. @@ -199,7 +183,9 @@ def test_no_web_fetch_tool_when_pill_off(monkeypatch): _drive(run()) tools = captured["body"].get("tools") or [] - assert all(t.get("type") != "web_fetch_20250910" for t in tools) + assert all( + t.get("type") not in ("web_fetch_20250910", "web_fetch_20260209") for t in tools + ) # ── SSE translation ───────────────────────────────────────────────── @@ -285,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. @@ -365,7 +356,9 @@ def test_web_fetch_error_renders_error_code(monkeypatch): def _finish_reasons(lines: list[str]) -> list: - """Return the finish_reason fields from every chat.completion.chunk.""" + """Return non-null finish_reason fields from each chat.completion.chunk. + Mid-stream content deltas carry ``finish_reason: None`` and are skipped + (the refusal path emits a notice delta before the content_filter chunk).""" out: list = [] for line in lines: if not line.startswith("data:"): @@ -380,8 +373,9 @@ def _finish_reasons(lines: list[str]) -> list: if parsed.get("object") != "chat.completion.chunk": continue for choice in parsed.get("choices") or []: - if "finish_reason" in choice: - out.append(choice["finish_reason"]) + reason = choice.get("finish_reason") + if reason is not None: + out.append(reason) return out 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_frontend_resolution.py b/studio/backend/tests/test_frontend_resolution.py new file mode 100644 index 0000000000..2b49763386 --- /dev/null +++ b/studio/backend/tests/test_frontend_resolution.py @@ -0,0 +1,248 @@ +# 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 the frontend-dist resolver in studio/backend/run.py. + +Loads only the relevant helpers via importlib so the test does not pull in +uvicorn / FastAPI / unsloth's full dependency tree. Pairs with the AST-style +test_host_defaults.py. +""" + +import ast +import importlib.util +import os +import sys +from pathlib import Path + +_RUN_PY = Path(__file__).resolve().parent.parent / "run.py" +_REPO_STUDIO_DIR = _RUN_PY.parent.parent # studio/ + + +def _load_helpers_only(): + """Import just the resolver helpers from run.py without executing the + server-side imports (uvicorn, structlog, etc.).""" + source = _RUN_PY.read_text(encoding = "utf-8") + tree = ast.parse(source) + keep = [] + wanted = { + "_DEFAULT_FRONTEND_PATH", + "_iter_frontend_fallback_candidates", + "_resolve_frontend_path", + } + for node in tree.body: + if isinstance(node, (ast.Import, ast.ImportFrom)): + keep.append(node) + elif isinstance(node, ast.Assign): + names = {t.id for t in node.targets if isinstance(t, ast.Name)} + if names & wanted: + keep.append(node) + elif isinstance(node, ast.FunctionDef) and node.name in wanted: + keep.append(node) + module = ast.Module(body = keep, type_ignores = []) + code = compile(module, str(_RUN_PY), "exec") + ns: dict = {"__file__": str(_RUN_PY), "__name__": "_run_helpers_test"} + exec(code, ns) + return ns + + +def test_resolver_returns_none_when_nothing_exists(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "no_studio")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](tmp_path / "missing") + assert chosen is None + assert attempted == [tmp_path / "missing"] + + +def test_resolver_picks_first_existing_candidate(tmp_path, monkeypatch): + dist = tmp_path / "good" / "frontend" / "dist" + dist.mkdir(parents = True) + (dist / "index.html").write_text("", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "no_studio")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](dist) + assert chosen == dist + assert attempted[-1] == dist + + +def test_resolver_falls_back_to_studio_home_site_packages(tmp_path, monkeypatch): + studio_home = tmp_path / "studio_home" + sp_dist = ( + studio_home + / "unsloth_studio" + / "lib" + / "python3.13" + / "site-packages" + / "studio" + / "frontend" + / "dist" + ) + sp_dist.mkdir(parents = True) + (sp_dist / "index.html").write_text("", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == sp_dist.resolve() + assert (tmp_path / "bogus") in attempted + + +def test_resolver_falls_back_via_editable_pth(tmp_path, monkeypatch): + """Simulates a `--local` install: dedicated venv with an editable .pth + pointing at a cloned repo that owns the built dist.""" + studio_home = tmp_path / "studio_home" + sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" + sp.mkdir(parents = True) + repo_root = tmp_path / "clone" + repo_studio = repo_root / "studio" + repo_dist = repo_studio / "frontend" / "dist" + repo_dist.mkdir(parents = True) + (repo_dist / "index.html").write_text("", encoding = "utf-8") + # Minimal `__editable___pkg_finder.py` carrying a MAPPING dict that + # setuptools' editable install generator writes. + finder = sp / "__editable___unsloth_0_0_0_finder.py" + finder.write_text( + "MAPPING: dict[str, str] = " + f"{{'studio': {str(repo_studio)!r}, 'unsloth': '/x', 'unsloth_cli': '/y'}}\n", + encoding = "utf-8", + ) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, attempted = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == repo_dist.resolve() + + +def test_iter_candidates_handles_missing_studio_home(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "nonexistent")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + # Glob over a non-existent dir is empty; must not raise. + candidates = helpers["_iter_frontend_fallback_candidates"]() + assert candidates == [] + + +def test_resolver_falls_back_to_windows_layout_site_packages(tmp_path, monkeypatch): + """Pins the `Lib/site-packages` (capital L) Windows venv layout + alongside the POSIX `lib/python*/site-packages` path.""" + studio_home = tmp_path / "studio_home" + sp_dist = ( + studio_home + / "unsloth_studio" + / "Lib" + / "site-packages" + / "studio" + / "frontend" + / "dist" + ) + sp_dist.mkdir(parents = True) + (sp_dist / "index.html").write_text("", encoding = "utf-8") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, _ = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == sp_dist.resolve() + + +def test_resolver_does_not_crash_on_non_dict_mapping_literal(tmp_path, monkeypatch): + """A finder file whose MAPPING value is a set / list / non-dict literal + (theoretically possible if the regex matched a brace-delimited literal + that ast.literal_eval can parse) must not AttributeError. The resolver + should skip that finder and keep probing.""" + studio_home = tmp_path / "studio_home" + sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" + sp.mkdir(parents = True) + # Bad finder: set literal, not a dict. ast.literal_eval parses it as set; + # any .get() call on it would raise AttributeError. + (sp / "__editable___bad_0_0_0_finder.py").write_text( + "MAPPING: dict[str, str] = {'studio', 'unsloth', 'unsloth_cli'}\n", + encoding = "utf-8", + ) + # Good finder that should still be discovered after the bad one is skipped. + repo_root = tmp_path / "clone" + repo_dist = repo_root / "studio" / "frontend" / "dist" + repo_dist.mkdir(parents = True) + (repo_dist / "index.html").write_text("", encoding = "utf-8") + (sp / "__editable___good_0_0_0_finder.py").write_text( + f"MAPPING: dict[str, str] = {{'studio': {str(repo_root / 'studio')!r}}}\n", + encoding = "utf-8", + ) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, _ = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == repo_dist.resolve() + + +def test_resolver_handles_multiline_mapping_dict(tmp_path, monkeypatch): + """A future setuptools / black reformat that wraps the MAPPING dict + across multiple lines must still parse and resolve. Locks in the + `[^}]*` + re.DOTALL behaviour.""" + studio_home = tmp_path / "studio_home" + sp = studio_home / "unsloth_studio" / "lib" / "python3.13" / "site-packages" + sp.mkdir(parents = True) + repo_root = tmp_path / "clone" + repo_studio = repo_root / "studio" + repo_dist = repo_studio / "frontend" / "dist" + repo_dist.mkdir(parents = True) + (repo_dist / "index.html").write_text("", encoding = "utf-8") + finder = sp / "__editable___unsloth_0_0_0_finder.py" + finder.write_text( + "MAPPING: dict[str, str] = {\n" + f" 'studio': {str(repo_studio)!r},\n" + " 'unsloth': '/x',\n" + " 'unsloth_cli': '/y',\n" + "}\n", + encoding = "utf-8", + ) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + chosen, _ = helpers["_resolve_frontend_path"](tmp_path / "bogus") + assert chosen is not None + assert chosen.resolve() == repo_dist.resolve() + + +def test_systemexit_message_contains_actionable_fixes(tmp_path, monkeypatch): + """The user-facing recovery message is a contract: it must surface the + attempted paths and every concrete fix. Pin its structure so a future + refactor doesn't drop one.""" + import os + import sys + + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "no_studio")) + monkeypatch.delenv("STUDIO_HOME", raising = False) + helpers = _load_helpers_only() + bogus = tmp_path / "no_such_dist" + _, attempted = helpers["_resolve_frontend_path"](bogus) + home = Path(os.environ["UNSLOTH_STUDIO_HOME"]).expanduser() + if sys.platform == "win32": + installer_bin = home / "bin" / "unsloth.exe" + else: + installer_bin = home / "unsloth_studio" / "bin" / "unsloth" + tried_lines = "\n".join(f" - {p}" for p in attempted) + message = ( + "[ERROR] Studio frontend build not found.\n" + f"Tried:\n{tried_lines}\n" + "\n" + "Likely cause: another 'unsloth' on PATH is shadowing the " + "installer's binary and points at a site-packages tree with " + "no built dist.\n" + "\n" + "Fix one of:\n" + f" - run the installer's binary directly: {installer_bin} studio\n" + " - pass --frontend \n" + " - pass --api-only to skip serving the web UI\n" + " - reinstall: curl -fsSL https://unsloth.ai/install.sh | sh" + ) + assert str(bogus) in message + assert "--frontend" in message + assert "--api-only" in message + assert "reinstall" in message + assert "installer's binary directly" in message + assert str(installer_bin) in message 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_index_bootstrap_origin.py b/studio/backend/tests/test_index_bootstrap_origin.py new file mode 100644 index 0000000000..89f7613ee4 --- /dev/null +++ b/studio/backend/tests/test_index_bootstrap_origin.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 + +"""Regression coverage for the bootstrap-pw cross-origin leak (PR 5739). +``_is_same_origin_request`` gates ``_inject_bootstrap`` so the seeded +admin password only ships to same-origin callers. +""" + +import os +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + + +def _build_request(host: str, origin: str | None, scheme: str = "http") -> MagicMock: + request = MagicMock() + request.url.scheme = scheme + request.url.netloc = host + request.headers = {"origin": origin} if origin is not None else {} + return request + + +def test_is_same_origin_request_missing_origin_is_same_origin(monkeypatch): + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = None) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_matching_origin_is_same_origin(): + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "http://127.0.0.1:8888") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_evil_origin_is_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "https://evil.example") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_scheme_mismatch_is_cross_origin(): + # https origin against an http listener is not same-origin. + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "https://127.0.0.1:8888") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_port_mismatch_is_cross_origin(): + # Same host different port is not same-origin per the web platform. + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8888", origin = "http://127.0.0.1:5173") + assert _is_same_origin_request(req) is False + + +# ── Canonicalisation: default-port stripping + case folding ───────── + + +def test_is_same_origin_request_https_default_port_stripped_on_origin(): + """RFC 6454 strips default ports on Origin; Starlette's netloc may still + carry ``:443``. Canonicalise both sides so this stays same-origin. + """ + from main import _is_same_origin_request + + req = _build_request( + "example.com:443", origin = "https://example.com", scheme = "https" + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_http_default_port_stripped_on_origin(): + from main import _is_same_origin_request + + req = _build_request("example.com:80", origin = "http://example.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_default_port_present_on_origin(): + """Mirror case: Origin carries the default port, netloc doesn't. Same-origin.""" + from main import _is_same_origin_request + + req = _build_request( + "example.com", origin = "https://example.com:443", scheme = "https" + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_host_case_insensitive(): + """Host portion is case-insensitive per RFC 3986.""" + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "http://EXAMPLE.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_scheme_case_insensitive(): + """Scheme portion is case-insensitive per RFC 3986.""" + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "HTTP://example.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_null_origin_is_cross_origin(): + """Sandboxed iframes / file:// pages send ``Origin: null``; cross-origin.""" + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "null") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_unparseable_origin_is_cross_origin(): + """Garbage values without a host fall to cross-origin; a malformed header + must not leak the bootstrap. + """ + from main import _is_same_origin_request + + req = _build_request("example.com", origin = "not-a-url") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_userinfo_in_netloc_ignored(): + """``user:pass@host:port`` netlocs (RFC 3986) must compare equal to the + credentials-less Origin. + """ + from main import _is_same_origin_request + + req = _build_request("user:pass@example.com:80", origin = "http://example.com") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_explicit_non_default_port_still_mismatch(): + """Canonicalisation does NOT collapse non-default ports to default.""" + from main import _is_same_origin_request + + req = _build_request( + "example.com", origin = "https://example.com:9999", scheme = "https" + ) + assert _is_same_origin_request(req) is False diff --git a/studio/backend/tests/test_index_bootstrap_origin_extra.py b/studio/backend/tests/test_index_bootstrap_origin_extra.py new file mode 100644 index 0000000000..aea6b36a96 --- /dev/null +++ b/studio/backend/tests/test_index_bootstrap_origin_extra.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Extra edge-case coverage for the bootstrap-pw cross-origin gate. +Companion to ``test_index_bootstrap_origin.py``: IPv6 netlocs, opaque +origins (``data:``, ``blob:``), comma-joined multi-Origin headers, and +the ``localhost`` vs ``127.0.0.1`` distinct-origin rule. +""" + +from unittest.mock import MagicMock + + +def _build_request(host: str, origin, scheme: str = "http") -> MagicMock: + request = MagicMock() + request.url.scheme = scheme + request.url.netloc = host + request.headers = {"origin": origin} if origin is not None else {} + return request + + +# ── IPv6 ──────────────────────────────────────────────────────────── + + +def test_is_same_origin_request_ipv6_loopback_same_origin(): + """Studio supports ``-H ::1`` binds; netloc is ``[::1]:8902``. Bare + ``partition(":")`` mis-parses the bracketed form and would refuse the + bootstrap on legitimate same-origin nav. + """ + from main import _is_same_origin_request + + req = _build_request("[::1]:8902", origin = "http://[::1]:8902") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_full_address_same_origin(): + from main import _is_same_origin_request + + req = _build_request( + "[2001:db8::1]:8443", + origin = "https://[2001:db8::1]:8443", + scheme = "https", + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_default_port_stripped(): + """Browser drops :80 on ``http://[::1]``.""" + from main import _is_same_origin_request + + req = _build_request("[::1]:80", origin = "http://[::1]") + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_case_insensitive(): + """Hex digits in IPv6 are case-insensitive per RFC 5952.""" + from main import _is_same_origin_request + + req = _build_request( + "[2001:DB8::1]:8443", + origin = "https://[2001:db8::1]:8443", + scheme = "https", + ) + assert _is_same_origin_request(req) is True + + +def test_is_same_origin_request_ipv6_different_host_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("[::1]:8902", origin = "http://[2001:db8::1]:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_ipv6_port_mismatch_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("[::1]:8902", origin = "http://[::1]:9999") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_ipv6_userinfo_stripped(): + from main import _is_same_origin_request + + req = _build_request("user:pass@[::1]:8902", origin = "http://[::1]:8902") + assert _is_same_origin_request(req) is True + + +# ── Opaque origins (data:, blob:) ─────────────────────────────────── + + +def test_is_same_origin_request_data_url_origin_is_cross_origin(): + """``data:`` URLs are opaque origins (HTML living standard); no host, + never same-origin. + """ + from main import _is_same_origin_request + + req = _build_request( + "127.0.0.1:8902", origin = "data:text/html," + ) + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_blob_url_origin_is_cross_origin(): + """``blob:`` URLs carry the inner origin only in non-canonical form; the + canonical comparison rejects them. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "blob:http://127.0.0.1:8902/uuid") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_file_url_origin_is_cross_origin(): + """``file://`` pages usually send ``Origin: null``; historical engines + sent ``Origin: file://``. Neither is same-origin vs an http listener. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "file://") + assert _is_same_origin_request(req) is False + + +# ── Multi-Origin header (comma-joined by Starlette) ──────────────── + + +def test_is_same_origin_request_comma_joined_origins_cross_origin(): + """Starlette concatenates repeated headers with ``, ``; the canonical + parser can't safely split this, so it falls to cross-origin. + """ + from main import _is_same_origin_request + + req = _build_request( + "127.0.0.1:8902", + origin = "http://127.0.0.1:8902, http://evil.example", + ) + assert _is_same_origin_request(req) is False + + +# ── localhost vs 127.0.0.1 (distinct origins per web platform) ────── + + +def test_is_same_origin_request_localhost_vs_127_is_cross_origin(): + """Browsers treat ``localhost`` and ``127.0.0.1`` as distinct origins; + the canonical comparison must not DNS-collapse them. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://localhost:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_127_vs_localhost_is_cross_origin(): + from main import _is_same_origin_request + + req = _build_request("localhost:8902", origin = "http://127.0.0.1:8902") + assert _is_same_origin_request(req) is False + + +# ── urlparse ValueError robustness ───────────────────────────────── + + +def test_is_same_origin_request_malformed_ipv6_bracket_is_cross_origin(): + """``urlparse`` raises ``ValueError('Invalid IPv6 URL')`` on unclosed + brackets (CVE-2024-11168 hardening). The gate must swallow and fall to + cross-origin rather than 500 the SPA handler. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://[malformed") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_invalid_ipv6_address_is_cross_origin(): + """Bracketed but invalid IPv6 (e.g. ``[::g]``) also raises + ``ValueError`` inside ``urlparse``.""" + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://[::g]:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_bracket_with_trailing_garbage_is_cross_origin(): + """Text after the closing bracket also raises inside ``urlparse``.""" + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "http://[2001:db8::1]extra:8902") + assert _is_same_origin_request(req) is False + + +def test_is_same_origin_request_empty_origin_header_is_cross_origin(): + """Explicit empty ``Origin:`` is not a valid serialised origin and must + not be conflated with a missing header; cross-origin, bootstrap withheld. + """ + from main import _is_same_origin_request + + req = _build_request("127.0.0.1:8902", origin = "") + assert _is_same_origin_request(req) is False diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 68a1c870fb..02a272ba3e 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -3,21 +3,35 @@ """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 +strip_shadowing_flags = _lsa.strip_shadowing_flags +validate_extra_args = _lsa.validate_extra_args # ── Pass-through (allowed) ─────────────────────────────────────────── @@ -60,13 +74,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 +112,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 +122,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 +161,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 +190,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 +260,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 +270,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 +282,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 +300,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 +331,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 +382,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", @@ -311,8 +411,7 @@ def test_is_managed_flag_false_for_mtp_pass_through(): 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 +442,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_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_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_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py index a431b78352..4d7528d238 100644 --- a/studio/backend/tests/test_multimodal_document.py +++ b/studio/backend/tests/test_multimodal_document.py @@ -117,6 +117,8 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch): types = [p.get("type") for p in parts] assert "document" in types, parts doc = _strip_cache(next(p for p in parts if p.get("type") == "document")) + # citations: {enabled: true} opts into Anthropic's natural-citation + # pipeline; without it the citations_delta handler is a no-op. assert doc == { "type": "document", "source": { @@ -124,6 +126,7 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch): "media_type": "application/pdf", "data": _TINY_PDF_B64, }, + "citations": {"enabled": True}, "title": "paper.pdf", } @@ -151,6 +154,7 @@ def test_anthropic_url_pdf_becomes_document_block(monkeypatch): assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, + "citations": {"enabled": True}, } @@ -255,6 +259,7 @@ def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch): assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, + "citations": {"enabled": True}, "title": "doc.pdf", } @@ -283,6 +288,7 @@ def test_anthropic_whitespace_only_data_uri_falls_back_to_file_url(monkeypatch): assert doc == { "type": "document", "source": {"type": "url", "url": "https://example.com/doc.pdf"}, + "citations": {"enabled": True}, } diff --git a/studio/backend/tests/test_openai_citation_markers.py b/studio/backend/tests/test_openai_citation_markers.py new file mode 100644 index 0000000000..ccc17be329 --- /dev/null +++ b/studio/backend/tests/test_openai_citation_markers.py @@ -0,0 +1,251 @@ +# 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 the OpenAI Responses-API citation marker rewriter. + +The stream interleaves text deltas with ``\\ue200cite\\ue202SOURCE_ID\\ue201`` +markers. The rewriter resolves each to `[N](URL)` when the annotation has +arrived and drops it otherwise; the URL list still flows to Sources via +`_record_url_citation`. + +Reference: https://developers.openai.com/api/docs/guides/citation-formatting +""" + +import pytest + +from core.inference.external_provider import ( + _replace_openai_citation_markers, + _rewrite_citation_markers_partial, +) + + +# Citation marker control codepoints (private-use area): +CITE_START = "" +CITE_STOP = "" +CITE_DELIM = "" + + +def _marker(source_id: str, locator: str | None = None) -> str: + payload = f"{CITE_START}cite{CITE_DELIM}{source_id}" + if locator: + payload = f"{payload}{CITE_DELIM}{locator}" + return f"{payload}{CITE_STOP}" + + +def _has_marker_codepoints(text: str) -> bool: + return any(c in text for c in (CITE_START, CITE_STOP, CITE_DELIM)) + + +def test_passthrough_when_no_marker_present(): + text = "Plain text with no citation markers." + assert _replace_openai_citation_markers(text, []) == text + + +def test_marker_rewritten_to_link_when_annotation_known(): + text = f"The capital is Paris {_marker('turn0view0')}." + citations = [ + { + "source_id": "turn0view0", + "url": "https://example.com/paris", + "title": "Paris", + }, + ] + out = _replace_openai_citation_markers(text, citations) + assert not _has_marker_codepoints(out) + assert "[[1]](https://example.com/paris)" in out + + +def test_unknown_source_marker_dropped_silently(): + text = f"Foo {_marker('turn9view9')} bar." + out = _replace_openai_citation_markers(text, []) + # Marker stripped, no garbled "E202" glyph leaks through, and the + # surrounding text stays intact. + assert not _has_marker_codepoints(out) + assert "E202" not in out + assert "turn9view9" not in out + assert "Foo" in out and "bar" in out + + +def test_multiple_concatenated_markers_resolved_in_order(): + """Real-world wire shape: a string of markers butted up against each other + after a sentence, as in the user-reported bug.""" + markers = "".join(_marker(f"turn{i}view{j}") for i, j in [(1, 0), (1, 1), (3, 0)]) + text = f"All animals ranked. {markers}" + citations = [ + {"source_id": "turn1view0", "url": "https://a.example/dog", "title": "Dog"}, + {"source_id": "turn1view1", "url": "https://a.example/cat", "title": "Cat"}, + {"source_id": "turn3view0", "url": "https://a.example/tiger", "title": "Tiger"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://a.example/dog)" in out + assert "[[2]](https://a.example/cat)" in out + assert "[[3]](https://a.example/tiger)" in out + assert not _has_marker_codepoints(out) + + +def test_marker_with_locator_resolves(): + text = f"See {_marker('turn2file0', 'L8-L13')}." + citations = [ + {"source_id": "turn2file0", "url": "https://example.com/doc.txt"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/doc.txt)" in out + assert "L8-L13" not in out # locator detail dropped; we just link. + assert not _has_marker_codepoints(out) + + +def test_mixed_known_and_unknown_markers(): + known = _marker("turn0view0") + unknown = _marker("turn0view99") + text = f"Known {known} and unknown {unknown}." + citations = [ + {"source_id": "turn0view0", "url": "https://example.com/known"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/known)" in out + # Unknown markers leave no trace, but surrounding prose stays. + assert "Known" in out and "unknown" in out + assert not _has_marker_codepoints(out) + assert "E202" not in out + + +def test_empty_text_returns_verbatim(): + assert _replace_openai_citation_markers("", []) == "" + + +def test_idempotent_on_pre_stripped_text(): + """Pre-stripped text (no private-use codepoints) returns verbatim.""" + text = "citeturn1view0 plain" + assert _replace_openai_citation_markers(text, []) == text + + +@pytest.mark.parametrize( + "citation", + [ + {"url": "https://example.com/a"}, # no source_id at all + {"source_id": None, "url": "https://example.com/b"}, + {"source_id": "", "url": "https://example.com/c"}, + ], +) +def test_citation_without_source_id_does_not_crash(citation): + text = f"X {_marker('turnXviewY')} Y" + out = _replace_openai_citation_markers(text, [citation]) + # No mapping, marker stripped. Crash-free is the contract. + assert not _has_marker_codepoints(out) + assert "turnXviewY" not in out + + +def test_multiple_source_id_aliases_resolve_to_same_url(): + """Every alias for the same URL must resolve, not just the first. + Regression for the Codex P1 on the original PR.""" + a = _marker("turn0view0") + b = _marker("turn0view0_span_1") + c = _marker("turn0view0_span_2") + text = f"Triple {a}{b}{c} cite." + citations = [ + { + "source_ids": ["turn0view0", "turn0view0_span_1", "turn0view0_span_2"], + "url": "https://example.com/paris", + "title": "Paris", + }, + ] + out = _replace_openai_citation_markers(text, citations) + # All three aliases collapse onto citation [1] -- the URL is the + # same so it would be misleading to show three different numbers. + assert out.count("[[1]](https://example.com/paris)") == 3 + assert not _has_marker_codepoints(out) + + +def test_source_ids_list_and_legacy_source_id_both_resolve(): + """Mixed-shape citation: legacy ``source_id`` plus newer + ``source_ids`` aliases both resolve.""" + legacy = _marker("legacy_id") + alias = _marker("alias_id") + text = f"Both {legacy} and {alias} work." + citations = [ + { + "source_id": "legacy_id", + "source_ids": ["alias_id"], + "url": "https://example.com/doc", + }, + ] + out = _replace_openai_citation_markers(text, citations) + assert out.count("[[1]](https://example.com/doc)") == 2 + assert not _has_marker_codepoints(out) + + +# --------------------------------------------------------------------------- +# _rewrite_citation_markers_partial: deferred-annotation tests. OpenAI emits +# url_citation annotations on a subsequent SSE event; this helper reports +# `has_unresolved` so the stream loop defers emission. See PR #5713 audit. +# --------------------------------------------------------------------------- + + +def test_partial_known_marker_resolves_and_clears_unresolved(): + text = f"Foo {_marker('s1')} bar." + out, unresolved = _rewrite_citation_markers_partial( + text, + [{"source_id": "s1", "url": "https://example.com/a"}], + ) + assert "[[1]](https://example.com/a)" in out + assert unresolved is False + assert not _has_marker_codepoints(out) + + +def test_partial_unknown_marker_preserves_verbatim_and_flags(): + text = f"Foo {_marker('s1')} bar." + out, unresolved = _rewrite_citation_markers_partial(text, []) + assert unresolved is True + # Codepoints must remain so a follow-up pass can re-parse. + assert _has_marker_codepoints(out) + assert "Foo" in out and "bar." in out + + +def test_partial_resolves_after_late_annotation(): + """Two-pass: first call sees no citations, second resolves after annotation.""" + text = f"See {_marker('s1')} for details." + out1, unresolved1 = _rewrite_citation_markers_partial(text, []) + assert unresolved1 is True + citations = [{"source_id": "s1", "url": "https://example.com/x"}] + out2, unresolved2 = _rewrite_citation_markers_partial(out1, citations) + assert unresolved2 is False + assert "[[1]](https://example.com/x)" in out2 + assert not _has_marker_codepoints(out2) + + +def test_partial_multi_source_partial_resolution_keeps_marker_pending(): + """Any unresolved token in a multi-source marker leaves the whole marker + verbatim with ``unresolved`` True; defer until every id resolves or + end-of-stream forces a flush (dropping unresolved tokens then).""" + cite = f"{CITE_START}cite{CITE_DELIM}known{CITE_DELIM}locator{CITE_STOP}" + text = f"Pre {cite} post." + citations = [{"source_id": "known", "url": "https://example.com/y"}] + out, unresolved = _rewrite_citation_markers_partial(text, citations) + assert unresolved is True + assert cite in out + # End-of-stream force flush: drop the unresolved token, keep the + # resolved link. The streamer routes pending segments through + # `_replace_openai_citation_markers` at force=True for this. + forced = _replace_openai_citation_markers(out, citations) + assert "[[1]](https://example.com/y)" in forced + assert "locator" not in forced + assert not _has_marker_codepoints(forced) + + +def test_partial_idempotent_on_marker_free_text(): + text = "Plain text." + out, unresolved = _rewrite_citation_markers_partial(text, []) + assert out == text + assert unresolved is False + + +def test_partial_mixed_known_and_pending_markers_flags_unresolved(): + known = _marker("known") + pending = _marker("pending") + text = f"{known} {pending}" + citations = [{"source_id": "known", "url": "https://example.com/k"}] + out, unresolved = _rewrite_citation_markers_partial(text, citations) + assert unresolved is True # the pending marker drives the flag + assert "[[1]](https://example.com/k)" in out + # The pending marker stays verbatim for the next pass. + assert CITE_START in out and "pending" in out diff --git a/studio/backend/tests/test_openai_citation_markers_edge.py b/studio/backend/tests/test_openai_citation_markers_edge.py new file mode 100644 index 0000000000..ffe8c6b6eb --- /dev/null +++ b/studio/backend/tests/test_openai_citation_markers_edge.py @@ -0,0 +1,413 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge-case tests for the OpenAI Responses citation marker rewriter. + +Covers multi-source markers, source+locator, marker SPLIT across SSE deltas, +unterminated tails at end-of-stream, multiple markers per delta, late +annotation ordering, and idempotency. + +Reference: https://developers.openai.com/api/docs/guides/citation-formatting +""" + +import importlib + + +# Streaming integration is exercised by ``_simulate_delta_stream`` further +# down, mirroring the head/buffer/flush dance from ``_stream_openai_responses``. +_module = importlib.import_module("core.inference.external_provider") +_replace_openai_citation_markers = _module._replace_openai_citation_markers +_split_pending_citation_tail = _module._split_pending_citation_tail + + +CITE_START = "" +CITE_STOP = "" +CITE_DELIM = "" + + +def _marker(*source_ids: str, locator: str | None = None) -> str: + """Build a ``\\ue200cite\\ue202[\\ue202...][\\ue202]\\ue201`` + marker. Accepts one or many ``source_ids`` plus an optional ``locator``.""" + payload = f"{CITE_START}cite{CITE_DELIM}" + CITE_DELIM.join(source_ids) + if locator: + payload = f"{payload}{CITE_DELIM}{locator}" + return f"{payload}{CITE_STOP}" + + +def _no_private_use(text: str) -> bool: + return all(c not in text for c in (CITE_START, CITE_STOP, CITE_DELIM)) + + +# Harness mirroring the head/pending-tail/flush dance in +# `_stream_openai_responses`, so streaming tests skip the httpx mock. +def _simulate_delta_stream( + deltas: list[str], + citations: list[dict], + *, + flush: bool = True, +) -> str: + pending = "" + emitted: list[str] = [] + for delta in deltas: + combined = pending + delta + head, pending = _split_pending_citation_tail(combined) + if head: + head = _replace_openai_citation_markers(head, citations) + if head: + emitted.append(head) + if flush and pending: + # Mirror `_flush_pending_marker_tail`: drop the tail entirely if no + # closing stop byte arrived; the literal ``cite`` would leak otherwise. + if CITE_STOP not in pending: + rendered = "" + else: + rendered = _replace_openai_citation_markers(pending, citations) + for ch in (CITE_START, CITE_STOP, CITE_DELIM): + rendered = rendered.replace(ch, "") + import re as _re + + rendered = _re.sub(r"^cite\S*", "", rendered) + if rendered: + emitted.append(rendered) + return "".join(emitted) + + +# --------------------------------------------------------------------------- +# 1. Multi-source markers per the OpenAI docs. +# --------------------------------------------------------------------------- + + +def test_multi_source_marker_all_resolve(): + """\\ue200cite\\ue202id1\\ue202id2\\ue202id3\\ue201 expands to three links + when every id is known. Earlier regex captured only id1 and dropped id2/id3.""" + text = f"All three: {_marker('id1', 'id2', 'id3')}" + citations = [ + {"source_id": "id1", "url": "https://example.com/1"}, + {"source_id": "id2", "url": "https://example.com/2"}, + {"source_id": "id3", "url": "https://example.com/3"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/1)" in out + assert "[[2]](https://example.com/2)" in out + assert "[[3]](https://example.com/3)" in out + assert _no_private_use(out) + + +def test_multi_source_marker_partial_resolution(): + """Known ids render, unknown ids drop silently, no glyph leaks.""" + text = f"Mixed: {_marker('known', 'unknown', 'also_known')}" + citations = [ + {"source_id": "known", "url": "https://k.example"}, + {"source_id": "also_known", "url": "https://ak.example"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://k.example)" in out + assert "[[2]](https://ak.example)" in out + assert "unknown" not in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 2. Source + locator: locator is dropped, link still resolves. +# --------------------------------------------------------------------------- + + +def test_marker_with_numeric_locator(): + text = f"See {_marker('tu0', locator = '42')}." + citations = [{"source_id": "tu0", "url": "https://example.com/doc"}] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/doc)" in out + assert "42" not in out + assert _no_private_use(out) + + +def test_marker_with_range_locator(): + text = f"See {_marker('tu0', locator = 'L8-L13')}." + citations = [{"source_id": "tu0", "url": "https://example.com/code"}] + out = _replace_openai_citation_markers(text, citations) + assert "[[1]](https://example.com/code)" in out + assert "L8-L13" not in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 3. Marker SPLIT across two SSE deltas -- the codex-flagged P1. +# --------------------------------------------------------------------------- + + +def test_marker_split_in_source_id(): + """Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 starts + with the rest (``rn0view0\\ue201``). The buffer stitches the halves + back together so they resolve to one link instead of leaking.""" + full = f"See {_marker('turn0view0')} now." + # Cut right after the second delim + "tu" inside the source id. + cut = full.index("tu", full.index(CITE_START)) + len("tu") + d1, d2 = full[:cut], full[cut:] + # Sanity check: delta-1 actually contains a partial marker. + assert CITE_START in d1 and CITE_STOP not in d1 + assert CITE_STOP in d2 + citations = [{"source_id": "turn0view0", "url": "https://x"}] + out = _simulate_delta_stream([d1, d2], citations) + assert out == "See [[1]](https://x) now." + assert _no_private_use(out) + + +def test_marker_split_at_start_byte(): + """Split exactly after the opening ``\\ue200`` byte; the buffer must + hold the lone open byte until the rest arrives.""" + full = f"Text {_marker('sid')} done" + cut = full.index(CITE_START) + 1 # right AFTER the open byte + d1, d2 = full[:cut], full[cut:] + citations = [{"source_id": "sid", "url": "https://y"}] + out = _simulate_delta_stream([d1, d2], citations) + assert out == "Text [[1]](https://y) done" + assert _no_private_use(out) + + +def test_marker_split_across_three_deltas(): + """Worst case: marker chopped into three pieces across three deltas.""" + full = f"A {_marker('threesplit')} B" + # cut at two points inside the marker + open_pos = full.index(CITE_START) + stop_pos = full.index(CITE_STOP) + cut1 = open_pos + 4 + cut2 = stop_pos - 2 + parts = [full[:cut1], full[cut1:cut2], full[cut2:]] + citations = [{"source_id": "threesplit", "url": "https://z"}] + out = _simulate_delta_stream(parts, citations) + assert out == "A [[1]](https://z) B" + assert _no_private_use(out) + + +def test_marker_split_with_trailing_text_after_close(): + """Delta-2 closes the marker AND carries trailing prose; both emit cleanly.""" + full = f"X {_marker('sid')} after" + cut = full.index("cite") + len("ci") + d1, d2 = full[:cut], full[cut:] + citations = [{"source_id": "sid", "url": "https://a"}] + out = _simulate_delta_stream([d1, d2], citations) + assert out == "X [[1]](https://a) after" + assert _no_private_use(out) + + +def test_split_marker_unknown_source_is_dropped_cleanly(): + """Split marker for an unknown source drops silently on flush.""" + full = f"Pre {_marker('never_seen')} post" + cut = full.index(CITE_START) + 3 + d1, d2 = full[:cut], full[cut:] + out = _simulate_delta_stream([d1, d2], []) + assert out == "Pre post" + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 4. Unterminated marker at end-of-stream -- truncation safety. +# --------------------------------------------------------------------------- + + +def test_unterminated_marker_at_stream_end_dropped_on_flush(): + """Stream ends mid-marker (e.g. response.incomplete); the tail is + flushed with private-use bytes stripped, no `E202` text leaks.""" + deltas = ["Some text ", f"{CITE_START}citetu", "rn0view0"] # no STOP ever + out = _simulate_delta_stream(deltas, [], flush = True) + assert _no_private_use(out) + assert "E200" not in out and "E202" not in out + # Surrounding prose stays; we don't assert exact marker remainder. + assert "Some text " in out + + +def test_flush_resolves_marker_when_late_annotation_arrives(): + """Marker in a delta, matching annotation arrives later (on + response.output_text.annotation.added after the final delta). The + rewriter reads ``all_url_citations`` LIVE at flush, so the buffered + marker still resolves.""" + deltas = ["Look ", f"{CITE_START}cite{CITE_DELIM}late_sid"] + pending = "" + citations: list[dict] = [] + emitted: list[str] = [] + for d in deltas: + combined = pending + d + head, pending = _split_pending_citation_tail(combined) + if head: + emitted.append(_replace_openai_citation_markers(head, citations)) + # Annotation arrives AFTER all deltas but BEFORE flush. + citations.append({"source_id": "late_sid", "url": "https://late.example"}) + # Append the STOP byte that closed the marker in a later delta. + pending = pending + CITE_STOP + flushed = _replace_openai_citation_markers(pending, citations) + for ch in (CITE_START, CITE_STOP, CITE_DELIM): + flushed = flushed.replace(ch, "") + emitted.append(flushed) + out = "".join(emitted) + assert "[[1]](https://late.example)" in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# 5. Multiple unrelated markers in a single delta. +# --------------------------------------------------------------------------- + + +def test_three_markers_in_one_delta_resolve_independently(): + text = f"alpha {_marker('a')} beta {_marker('b')} gamma {_marker('c')} end" + citations = [ + {"source_id": "a", "url": "https://example.com/a"}, + {"source_id": "b", "url": "https://example.com/b"}, + {"source_id": "c", "url": "https://example.com/c"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert out == ( + "alpha [[1]](https://example.com/a) beta " + "[[2]](https://example.com/b) gamma " + "[[3]](https://example.com/c) end" + ) + + +# --------------------------------------------------------------------------- +# 6. Idempotency. +# --------------------------------------------------------------------------- + + +def test_rewriter_idempotent_on_already_rewritten_text(): + """Running the rewriter twice does not double-link or corrupt brackets.""" + text = f"alpha {_marker('a')} omega" + citations = [{"source_id": "a", "url": "https://example.com/a"}] + once = _replace_openai_citation_markers(text, citations) + twice = _replace_openai_citation_markers(once, citations) + assert once == twice + assert _no_private_use(once) + + +def test_rewriter_idempotent_on_marker_free_text(): + """No-op when there is nothing to rewrite.""" + text = "Plain prose with no citations and no private-use bytes." + out = _replace_openai_citation_markers(text, []) + assert out is text or out == text + + +# --------------------------------------------------------------------------- +# 7. Edge / robustness. +# --------------------------------------------------------------------------- + + +def test_only_marker_no_surrounding_text(): + """A delta that is JUST a marker (no prose) still renders correctly; + used to leak without the empty-string short-circuit in the split helper.""" + text = _marker("solo") + citations = [{"source_id": "solo", "url": "https://solo.example"}] + out = _replace_openai_citation_markers(text, citations) + assert out == "[[1]](https://solo.example)" + + +def test_back_to_back_markers_with_no_separator(): + """Adjacent markers resolve to concatenated links, no joining whitespace.""" + text = f"{_marker('x')}{_marker('y')}" + citations = [ + {"source_id": "x", "url": "https://x.example"}, + {"source_id": "y", "url": "https://y.example"}, + ] + out = _replace_openai_citation_markers(text, citations) + assert out == "[[1]](https://x.example)[[2]](https://y.example)" + + +def test_split_helper_buffers_only_after_last_open_byte(): + """A complete marker followed by an unterminated one: head includes + the complete marker, buffer holds only the trailing partial.""" + complete = _marker("done") + partial = f"{CITE_START}cite{CITE_DELIM}half" # no STOP + text = f"pre {complete} mid {partial}" + head, tail = _split_pending_citation_tail(text) + assert head == f"pre {complete} mid " + assert tail == partial + # And the head, once rewritten, drops every private-use byte. + rewritten = _replace_openai_citation_markers( + head, [{"source_id": "done", "url": "https://d"}] + ) + assert rewritten == "pre [[1]](https://d) mid " + + +def test_split_helper_empty_input(): + head, tail = _split_pending_citation_tail("") + assert head == "" and tail == "" + + +def test_split_helper_no_open_byte(): + head, tail = _split_pending_citation_tail("nothing to see here") + assert head == "nothing to see here" and tail == "" + + +def test_split_helper_complete_marker_only(): + """A delta ending with a closed marker leaves the buffer empty.""" + text = f"alpha {_marker('a')}" + head, tail = _split_pending_citation_tail(text) + assert head == text and tail == "" + + +# --------------------------------------------------------------------------- +# 8. Sources-panel: marker drop must not affect citation aggregation. +# Indices come from the url_citations list, not the marker stream. +# --------------------------------------------------------------------------- + + +def test_unknown_marker_does_not_perturb_citation_indexing(): + """Unknown source_id markers drop without consuming an index slot.""" + text = f"A {_marker('unknown')} B {_marker('real_a')} C {_marker('real_b')}" + citations = [ + {"source_id": "real_a", "url": "https://example.com/a"}, + {"source_id": "real_b", "url": "https://example.com/b"}, + ] + out = _replace_openai_citation_markers(text, citations) + # real_a is index 1; unknown does not take a slot. + assert "[[1]](https://example.com/a)" in out + assert "[[2]](https://example.com/b)" in out + assert _no_private_use(out) + + +# --------------------------------------------------------------------------- +# Regression: unterminated marker tail must NOT leak the residual +# ``cite``-prefixed source id as plain text. PR #5713 audit P1. +# --------------------------------------------------------------------------- + + +def test_unterminated_marker_does_not_leak_cite_residue(): + """Stream ends mid-marker: drop the whole tail rather than strip + codepoints and leave ``cite`` behind.""" + half = f"Hi there {CITE_START}cite{CITE_DELIM}turn0view0" + out = _simulate_delta_stream([half], [], flush = True) + # Prose before the marker stays; no private-use bytes or cite residue. + assert "Hi there" in out + assert _no_private_use(out) + assert "citeturn0view0" not in out + assert "cite" not in out.split("Hi there", 1)[1] + + +def test_unterminated_marker_only_no_prefix_drops_entirely(): + """A delta that is purely an unterminated marker flushes to "".""" + half = f"{CITE_START}cite{CITE_DELIM}turn0view0" + out = _simulate_delta_stream([half], [], flush = True) + assert out == "" + + +def test_unterminated_marker_with_prefix_emits_only_prefix(): + """Prose then unterminated marker: prose emits, marker remnant drops.""" + half = f"prefix prose {CITE_START}cite{CITE_DELIM}abc" + out = _simulate_delta_stream([half], [], flush = True) + assert out == "prefix prose " + + +def test_closing_byte_arrives_after_pending_buffered_split(): + """Closing byte arrives in a later delta after opener + source id were + buffered; link resolves with no residue.""" + cuts = [ + f"a {CITE_START}cite{CITE_DELIM}", + f"sid{CITE_STOP} b", + ] + out = _simulate_delta_stream( + cuts, + [{"source_id": "sid", "url": "https://example.com/x"}], + flush = True, + ) + assert "[[1]](https://example.com/x)" in out + assert "a " in out and "b" in out + assert _no_private_use(out) + assert "citesid" not in out 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 1f9c2710e4..f5dee2561a 100644 --- a/studio/backend/tests/test_openai_image_generation.py +++ b/studio/backend/tests/test_openai_image_generation.py @@ -207,9 +207,13 @@ 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" assert ends[0]["image_mime"] == "image/png" 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_openai_tool_result_fallbacks.py b/studio/backend/tests/test_openai_tool_result_fallbacks.py new file mode 100644 index 0000000000..7c033bc348 --- /dev/null +++ b/studio/backend/tests/test_openai_tool_result_fallbacks.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for OpenAI Responses tool-result rendering. + +Covers two bug classes: empty web_search cards (per-card result seeded +with "Searching: ") and orphan shell_call cards (bundled-output +fallback + final flush at response.completed / response.incomplete). +""" + +import asyncio +import json + +import httpx + +from core.inference import external_provider as ep_mod +from core.inference.external_provider import ExternalProviderClient + + +def _drive(coro): + return asyncio.new_event_loop().run_until_complete(coro) + + +async def _collect(agen): + out = [] + async for line in agen: + out.append(line) + return out + + +def _mock_http_client(monkeypatch, handler): + transport = httpx.MockTransport(handler) + monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport)) + + +def _make_client(base_url: str = "https://api.openai.com/v1") -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "openai", + base_url = base_url, + api_key = "sk-test", + ) + + +def _openai_sse(events: list[dict]) -> bytes: + chunks: list[str] = [] + for event in events: + chunks.append(f"event: {event['type']}") + chunks.append(f"data: {json.dumps(event)}") + chunks.append("") + return ("\n".join(chunks) + "\n").encode("utf-8") + + +def _tool_events(lines: list[str]) -> list[dict]: + out: list[dict] = [] + for line in lines: + if not line.startswith("data:"): + continue + raw = line[len("data:") :].strip() + if not raw or raw == "[DONE]": + continue + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict) and "_toolEvent" in parsed: + out.append(parsed["_toolEvent"]) + return out + + +def _drive_stream(sse_events, enabled_tools, monkeypatch): + def handler(request): + return httpx.Response( + 200, + content = _openai_sse(sse_events), + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_client() + return await _collect( + client._stream_openai_responses( + messages = [{"role": "user", "content": "x"}], + model = "gpt-5.5", + temperature = 0.7, + top_p = 0.95, + max_tokens = 4096, + enable_thinking = None, + reasoning_effort = None, + enabled_tools = enabled_tools, + ) + ) + + return _drive(run()) + + +# ── web_search per-card result ───────────────────────────────────────── + + +def test_web_search_each_call_carries_its_own_query_as_result(monkeypatch): + """Each card carries its own `Searching: ` text; no empties.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_1", + "action": {"query": "popular animals 2026"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_2", + "action": {"query": "most loved animals poll"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_3", + "action": {"query": "tiger ranking"}, + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["web_search"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + by_id = {e["tool_call_id"]: e for e in ends} + assert by_id["ws_1"]["result"] == "Searching: popular animals 2026" + assert by_id["ws_2"]["result"] == "Searching: most loved animals poll" + assert by_id["ws_3"]["result"] == "Searching: tiger ranking" + + +def test_web_search_last_call_overwritten_with_citations(monkeypatch): + """Last call still gets the aggregated citation list; earlier calls + keep their per-call `Searching:` text.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_1", + "action": {"query": "first query"}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_2", + "action": {"query": "second query"}, + }, + }, + { + "type": "response.output_text.annotation.added", + "annotation": { + "type": "url_citation", + "url": "https://example.com/a", + "title": "Example A", + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["web_search"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + by_id: dict = {} + # Keep the LAST tool_end per id (the citation overwrite for ws_2). + for e in ends: + by_id[e["tool_call_id"]] = e + # First call keeps its own query. + assert by_id["ws_1"]["result"] == "Searching: first query" + # Last call gets overwritten with the citation block. + assert "Title: Example A" in by_id["ws_2"]["result"] + assert "URL: https://example.com/a" in by_id["ws_2"]["result"] + + +def test_web_search_empty_query_falls_back_to_empty_result(monkeypatch): + """No query -> empty result (no `Searching:` placeholder).""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "id": "ws_only", + "action": {}, + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["web_search"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert len(ends) == 1 + assert ends[0]["result"] == "" + + +# ── shell_call output fallbacks ──────────────────────────────────────── + + +def test_shell_call_emits_tool_end_when_output_bundled_on_done(monkeypatch): + """Output bundled on the shell_call done event emits tool_end.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_bundled", + "action": {"commands": ["echo hi"]}, + "output": [ + { + "stdout": "hi\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + starts = [e for e in events if e["type"] == "tool_start"] + ends = [e for e in events if e["type"] == "tool_end"] + assert len(starts) == 1 + assert starts[0]["tool_call_id"] == "scall_bundled" + assert len(ends) == 1 + assert ends[0]["tool_call_id"] == "scall_bundled" + assert "hi" in ends[0]["result"] + + +def test_shell_call_bundled_then_separate_output_does_not_double_emit(monkeypatch): + """Separate shell_call_output after bundled-output is a no-op.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_both", + "action": {"commands": ["echo bundle"]}, + "output": [ + { + "stdout": "bundle\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call_output", + "id": "scout_both", + "call_id": "scall_both", + "output": [ + { + "stdout": "should not double-emit\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert len(ends) == 1 + assert ends[0]["tool_call_id"] == "scall_both" + assert "bundle" in ends[0]["result"] + assert "should not double-emit" not in ends[0]["result"] + + +def test_shell_call_final_flush_on_completed_when_no_output_event(monkeypatch): + """Orphan shell_call finalises via the response.completed flush.""" + sse_events = [ + { + "type": "response.output_item.added", + "item": { + "type": "shell_call", + "id": "scall_orphan", + "action": {"commands": ["true"]}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_orphan", + "action": {"commands": ["true"]}, + "status": "completed", + }, + }, + {"type": "response.completed", "response": {}}, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert any(e["tool_call_id"] == "scall_orphan" for e in ends) + + +def test_shell_call_flushed_on_response_incomplete_truncation(monkeypatch): + """Truncated streams (response.incomplete) also flush orphan calls.""" + sse_events = [ + { + "type": "response.output_item.added", + "item": { + "type": "shell_call", + "id": "scall_truncated", + "action": {"commands": ["long_running"]}, + }, + }, + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_truncated", + "action": {"commands": ["long_running"]}, + "status": "in_progress", + }, + }, + { + "type": "response.incomplete", + "response": { + "incomplete_details": {"reason": "max_output_tokens"}, + }, + }, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert any(e["tool_call_id"] == "scall_truncated" for e in ends) + + +def test_shell_call_incomplete_does_not_double_emit(monkeypatch): + """response.incomplete is idempotent against already-finalised calls.""" + sse_events = [ + { + "type": "response.output_item.done", + "item": { + "type": "shell_call", + "id": "scall_done", + "action": {"commands": ["echo done"]}, + "output": [ + { + "stdout": "done\n", + "stderr": "", + "outcome": {"type": "exit", "exit_code": 0}, + } + ], + }, + }, + { + "type": "response.incomplete", + "response": { + "incomplete_details": {"reason": "max_output_tokens"}, + }, + }, + ] + lines = _drive_stream(sse_events, ["code_execution"], monkeypatch) + events = _tool_events(lines) + ends = [e for e in events if e["type"] == "tool_end"] + assert len(ends) == 1 + assert ends[0]["tool_call_id"] == "scall_done" + assert "done" in ends[0]["result"] diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py index cc8c16993c..313c15a441 100644 --- a/studio/backend/tests/test_pricing.py +++ b/studio/backend/tests/test_pricing.py @@ -1,12 +1,8 @@ # 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 per-session cost calculator. - -Pricing inputs are baked into ``core/inference/pricing.py``; this -test verifies the math (with multipliers from the prompt-caching -docs) and that unknown models / empty usage degrade gracefully. -""" +"""Unit tests for the per-session cost calculator. Verifies math +against ``core/inference/pricing.py`` and graceful degradation.""" import math @@ -14,6 +10,7 @@ from core.inference.pricing import ( ANTHROPIC_CACHE_5M_WRITE_MULT, ANTHROPIC_CACHE_1H_WRITE_MULT, ANTHROPIC_CACHE_READ_MULT, + ANTHROPIC_FAST_MODE_MULT, ANTHROPIC_PRICING, OPENAI_CACHE_READ_MULT, OPENAI_CONTAINER_USD_PER_HOUR, @@ -57,6 +54,64 @@ def test_anthropic_opus_4_7_input_and_output_math(): assert _isclose(out["total_usd"], 30.0) +# ── Anthropic fast-mode 6x multiplier (Opus 4.6 / 4.7 only) ───────── + + +def test_anthropic_fast_mode_charges_6x_standard_opus(): + """6x on input + output when ``usage.speed == "fast"``. + https://platform.claude.com/docs/en/build-with-claude/fast-mode""" + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 1_000_000, + "output_tokens": 1_000_000, + "speed": "fast", + }, + ) + assert _isclose(out["input_usd"], 5.0 * ANTHROPIC_FAST_MODE_MULT) + assert _isclose(out["output_usd"], 25.0 * ANTHROPIC_FAST_MODE_MULT) + assert _isclose(out["total_usd"], 30.0 * ANTHROPIC_FAST_MODE_MULT) + assert "(fast)" in out["model_priced"], out["model_priced"] + + +def test_anthropic_fast_mode_does_not_affect_standard_speed(): + """``speed: "standard"`` (or missing) keeps the base rates.""" + out_standard = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 1_000_000, + "output_tokens": 1_000_000, + "speed": "standard", + }, + ) + out_missing = calculate_cost( + "anthropic", + "claude-opus-4-7", + {"input_tokens": 1_000_000, "output_tokens": 1_000_000}, + ) + assert _isclose(out_standard["total_usd"], out_missing["total_usd"]) + assert _isclose(out_standard["total_usd"], 30.0) + + +def test_anthropic_fast_mode_stacks_with_cache_read_multiplier(): + """Cache multipliers apply on top of fast-mode (per docs).""" + base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 0, + "output_tokens": 0, + "cache_read_input_tokens": 1_000_000, + "speed": "fast", + }, + ) + expected = base * ANTHROPIC_FAST_MODE_MULT * ANTHROPIC_CACHE_READ_MULT + assert _isclose(out["cache_read_usd"], expected) + + # ── Anthropic cache write 5m + read multipliers ────────────────────── @@ -102,8 +157,7 @@ def test_anthropic_cache_1h_write_uses_2x_multiplier(): def test_anthropic_cache_5m_default_when_no_breakdown(): - # When the docs/response doesn't surface the 5m/1h split, treat - # the full cache_creation bucket as 5m (the upstream default pool). + # No 5m/1h split surfaced -> assume the default 5m pool. base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] out = calculate_cost( "anthropic", @@ -148,8 +202,7 @@ def test_anthropic_code_exec_charged_per_hour(): def test_anthropic_dated_id_falls_back_to_canonical_prefix(): - # Hypothetical dated snapshot of claude-opus-4-7 should still - # inherit the canonical-id pricing via the prefix-match fallback. + # Dated snapshot inherits canonical pricing via prefix-match. out = calculate_cost( "anthropic", "claude-opus-4-7-20260712", @@ -163,8 +216,7 @@ def test_anthropic_dated_id_falls_back_to_canonical_prefix(): def test_openai_gpt55_input_output_math(): - # Sub-272k input keeps us in the short-context tier ($5/$30). - # The dedicated long-context tests below exercise the crossover. + # Sub-272k stays in short-context tier ($5/$30). out = calculate_cost( "openai", "gpt-5.5", @@ -176,11 +228,7 @@ def test_openai_gpt55_input_output_math(): def test_openai_cache_read_subtracted_from_input_at_discount(): - # OpenAI folds cached tokens into input_tokens, unlike Anthropic. - # The calculator must subtract cached_tokens from the "full price" - # bucket and re-bill them at 0.1x. Use a sub-272k total so the - # short-context tier applies (long-context crossover is exercised - # in its own test below). + # OpenAI folds cached into input_tokens; subtract and re-bill at 0.1x. base = OPENAI_PRICING["gpt-5.5"]["input_per_mtok"] out = calculate_cost( "openai", @@ -199,9 +247,7 @@ def test_openai_cache_read_subtracted_from_input_at_discount(): def test_openai_billable_input_tokens_does_not_double_count_cache_read(): - # OpenAI's input_tokens already includes cached_tokens, so the - # billable counter must NOT add cache_read on top -- otherwise the - # tooltip says 180k input when the bill is for 100k. + # input_tokens already includes cached; don't double-count. out = calculate_cost( "openai", "gpt-5.5", @@ -215,9 +261,7 @@ def test_openai_billable_input_tokens_does_not_double_count_cache_read(): def test_openai_dated_snapshot_inherits_canonical_pricing(): - # Sub-272k stays in the short-context tier; the prefix-match - # fallback is what proves the dated snapshot inherits gpt-5.5 - # pricing. + # Dated snapshot inherits gpt-5.5 pricing via prefix-match. out = calculate_cost( "openai", "gpt-5.5-2026-04-23", @@ -228,10 +272,7 @@ def test_openai_dated_snapshot_inherits_canonical_pricing(): def test_openai_gpt54_family_uses_verified_prices(): - # Spot-check the lower-tier rows that previously underbilled. - # gpt-5.4 has a long-context tier so the input has to stay - # below 272k; the mini/nano/codex rows have no crossover so - # 1M tokens is fine. + # Spot-check lower-tier rows that previously underbilled. cases = { # (input_tokens, expected_input_usd, expected_output_usd) "gpt-5.4": (200_000, 200_000 / 1_000_000.0 * 2.5, 200_000 / 1_000_000.0 * 15.0), @@ -251,8 +292,7 @@ def test_openai_gpt54_family_uses_verified_prices(): def test_openai_unlisted_model_priced_false_not_zero_default(): - # o-series / gpt-4.5 are no longer on the pricing page, so we - # intentionally drop them rather than silently underbill at $0. + # o-series / gpt-4.5 are off the pricing page; drop rather than $0. for model in ("o3", "o4-mini", "gpt-4.5", "gpt-4.5-preview"): out = calculate_cost( "openai", @@ -270,9 +310,7 @@ def test_openai_unlisted_model_priced_false_not_zero_default(): def test_anthropic_canonical_4_5_ids_are_priced(): - # Codex P1: claude-opus-4-5 (no date) is the canonical id used - # in backend defaults but was missing from the table, so the - # calculator returned priced=False + zero cost. Pin the aliases. + # Pin the bare-id aliases (backend defaults reference these). cases = { "claude-opus-4-5": (5.0, 25.0), "claude-sonnet-4-5": (3.0, 15.0), @@ -307,8 +345,7 @@ def test_openai_gpt55_short_context_under_272k_uses_base_rates(): def test_openai_gpt55_long_context_crossover_uses_higher_rates(): - # 300k billable input > 272k threshold -> long-context tier - # applies to the WHOLE turn, not a per-token blend. + # >272k billable -> long-context tier on the whole turn. out = calculate_cost( "openai", "gpt-5.5", @@ -330,8 +367,7 @@ def test_openai_gpt54_long_context_crossover(): def test_openai_gpt54_mini_has_no_long_context_tier(): - # Mini/nano/codex don't publish a long-context price; the base - # rate must keep applying even at very large prompts. + # Mini/nano/codex have no long-context tier; base rate always applies. out = calculate_cost( "openai", "gpt-5.4-mini", @@ -374,8 +410,7 @@ def test_openai_container_hours_charged(): def test_openai_tool_surcharges_added_to_total(): - # End-to-end: input + output + web_search + container in one - # turn. Total must sum all four buckets. + # End-to-end: total must sum input + output + web_search + container. out = calculate_cost( "openai", "gpt-5.5", @@ -412,12 +447,12 @@ def test_snapshot_contains_provider_buckets_and_multipliers(): assert a["cache_5m_write_mult"] == ANTHROPIC_CACHE_5M_WRITE_MULT assert a["cache_1h_write_mult"] == ANTHROPIC_CACHE_1H_WRITE_MULT assert a["cache_read_mult"] == ANTHROPIC_CACHE_READ_MULT + assert a["fast_mode_mult"] == ANTHROPIC_FAST_MODE_MULT assert "web_search_usd_per_1k" in a assert "code_execution_usd_per_hour" in a assert "models" in o and "gpt-5.5" in o["models"] assert o["cache_read_mult"] == OPENAI_CACHE_READ_MULT - # OpenAI tool surcharge constants are also exposed so the frontend - # tooltip can render the per-call rate. + # OpenAI tool surcharge constants are exposed for the frontend. assert o["web_search_usd_per_1k"] == OPENAI_WEB_SEARCH_USD_PER_1K assert o["container_usd_per_hour"] == OPENAI_CONTAINER_USD_PER_HOUR # Long-context tier metadata travels with the model row. @@ -425,3 +460,169 @@ def test_snapshot_contains_provider_buckets_and_multipliers(): assert gpt55["long_context_threshold"] == 272_000 assert gpt55["long_context_input_per_mtok"] == 10.0 assert gpt55["long_context_output_per_mtok"] == 45.0 + + +# ── longest-prefix match: dated mini variant must not collide with the +# shorter family prefix. ── + + +def test_longest_prefix_match_wins_for_dated_mini_snapshot(): + """`gpt-5.4-mini-2026-...` must inherit the mini rate, not the + shorter `gpt-5.4` rate (longest prefix wins).""" + out = calculate_cost( + "openai", + "gpt-5.4-mini-2026-04-23", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is True + # mini = 0.75/MTok, shorter gpt-5.4 = 2.5/MTok (>3x overcharge). + assert _isclose(out["input_usd"], 0.75), out + + +def test_longest_prefix_match_wins_for_dated_pro_snapshot(): + out = calculate_cost( + "openai", + "gpt-5.5-pro-2026-04-23", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is True + # gpt-5.5-pro = 30/MTok vs gpt-5.5 = 5/MTok; longest wins. + assert _isclose(out["input_usd"], 30.0), out + + +# ── accept both chat-style and Responses envelope shapes. ── + + +def test_openai_chat_style_usage_keys_priced_correctly(): + """Chat-style envelope (`prompt_tokens` / `completion_tokens`) must + produce a non-zero cost (previously silently zeroed).""" + out = calculate_cost( + "openai", + "gpt-5.4-mini", + {"prompt_tokens": 1_000_000, "completion_tokens": 1_000_000}, + ) + # gpt-5.4-mini: 0.75 input + 4.5 output per MTok. + assert _isclose(out["input_usd"], 0.75), out + assert _isclose(out["output_usd"], 4.5), out + + +def test_input_tokens_preferred_when_both_keys_present(): + """Raw key wins when both envelope shapes are present.""" + out = calculate_cost( + "openai", + "gpt-5.4-mini", + { + "input_tokens": 2_000_000, + "prompt_tokens": 5_000_000, + "output_tokens": 0, + }, + ) + # input_tokens=2M wins -> 2 * 0.75 = 1.50. + assert _isclose(out["input_usd"], 1.50), out + + +def test_anthropic_chat_style_prompt_tokens_dedupes_cache_buckets(): + """Anthropic chat-style prompt_tokens already folds cache buckets; + don't double-count billable input.""" + # 1M uncached + 200K cache_creation + 500K cache_read -> 1.7M folded. + raw = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 1_000_000, + "cache_creation_input_tokens": 200_000, + "cache_read_input_tokens": 500_000, + "output_tokens": 0, + }, + ) + chat = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "prompt_tokens": 1_700_000, + "cache_creation_input_tokens": 200_000, + "cache_read_input_tokens": 500_000, + "completion_tokens": 0, + }, + ) + # Both envelopes must price the same. + assert _isclose(chat["input_usd"], raw["input_usd"]), (chat, raw) + assert _isclose(chat["cache_write_usd"], raw["cache_write_usd"]), (chat, raw) + assert _isclose(chat["cache_read_usd"], raw["cache_read_usd"]), (chat, raw) + assert _isclose(chat["total_usd"], raw["total_usd"]), (chat, raw) + assert chat["billable_input_tokens"] == raw["billable_input_tokens"], (chat, raw) + + +def test_openai_chat_style_prompt_tokens_keeps_cache_read_semantics(): + """OpenAI prompt_tokens includes cache_read like raw input_tokens.""" + raw = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": 1_000_000, + "input_tokens_details": {"cached_tokens": 200_000}, + "output_tokens": 100_000, + }, + ) + chat = calculate_cost( + "openai", + "gpt-5.5", + { + "prompt_tokens": 1_000_000, + "cache_read_input_tokens": 200_000, + "completion_tokens": 100_000, + }, + ) + assert _isclose(chat["total_usd"], raw["total_usd"]), (chat, raw) + + +def test_openai_chat_style_envelope_reads_cache_from_prompt_tokens_details(): + """Chat-style envelope ships cached under prompt_tokens_details; + calculator must honour both this and input_tokens_details.""" + base = OPENAI_PRICING["gpt-5.5"]["input_per_mtok"] + raw = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": 100_000, + "input_tokens_details": {"cached_tokens": 80_000}, + "output_tokens": 0, + }, + ) + chat_style = calculate_cost( + "openai", + "gpt-5.5", + { + "prompt_tokens": 100_000, + "prompt_tokens_details": {"cached_tokens": 80_000}, + "completion_tokens": 0, + }, + ) + # Both envelopes must price identically. + assert _isclose(chat_style["input_usd"], raw["input_usd"]), (chat_style, raw) + assert _isclose(chat_style["cache_read_usd"], raw["cache_read_usd"]), ( + chat_style, + raw, + ) + # 80k at 0.1x base, 20k at full. + assert _isclose( + chat_style["cache_read_usd"], + 80_000 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT, + ) + + +def test_explicit_zero_output_tokens_wins_over_stale_completion_tokens(): + """Explicit ``output_tokens: 0`` beats a stale ``completion_tokens``; + the previous `or` fallback treated 0 as missing.""" + out = calculate_cost( + "openai", + "gpt-4o-mini", + { + "input_tokens": 100, + "output_tokens": 0, + # Stale chat-style mirror; must not bill against it. + "completion_tokens": 50, + }, + ) + assert out["billable_output_tokens"] == 0, out + assert out["output_usd"] == 0.0, out diff --git a/studio/backend/tests/test_pricing_edge.py b/studio/backend/tests/test_pricing_edge.py new file mode 100644 index 0000000000..ca4be258e0 --- /dev/null +++ b/studio/backend/tests/test_pricing_edge.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Adversarial edge cases for ``calculate_cost`` / ``_lookup``: prefix +boundary, negative tokens, chat vs raw parity, long-context crossover +on billable count, and malformed sub-objects.""" + +import math + +from core.inference.pricing import ( + ANTHROPIC_CACHE_5M_WRITE_MULT, + ANTHROPIC_CACHE_READ_MULT, + ANTHROPIC_PRICING, + OPENAI_CACHE_READ_MULT, + OPENAI_PRICING, + _lookup, + calculate_cost, +) + + +def _isclose(a, b, tol = 1e-6): + return math.isclose(a, b, rel_tol = tol, abs_tol = tol) + + +# ── prefix-match boundary checks ──────────────────────────────────── + + +def test_prefix_match_requires_dash_boundary_opus_variant(): + # `claude-opus-4-15` must not inherit `claude-opus-4-1` pricing; + # next char must be `-` or end-of-string. + assert _lookup("anthropic", "claude-opus-4-15") is None + out = calculate_cost( + "anthropic", + "claude-opus-4-15", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is False + assert out["total_usd"] == 0.0 + + +def test_prefix_match_requires_dash_boundary_gpt_variant(): + # Same dash-boundary invariant for OpenAI ids. + assert _lookup("openai", "gpt-5.55") is None + assert _lookup("openai", "gpt-5.55-2026-04-23") is None + out = calculate_cost( + "openai", + "gpt-5.55-2026-04-23", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is False + + +def test_prefix_match_requires_dash_boundary_pro_lookalike(): + # `gpt-5.5-prod` must fall through `gpt-5.5-pro` (6x overcharge) + # and land on the canonical `gpt-5.5` row. + prices = _lookup("openai", "gpt-5.5-prod") + assert prices is not None + assert ( + prices["input_per_mtok"] == OPENAI_PRICING["gpt-5.5"]["input_per_mtok"] + ), "expected fallback to gpt-5.5 base ($5), not gpt-5.5-pro ($30)" + out = calculate_cost( + "openai", + "gpt-5.5-prod", + {"input_tokens": 100_000, "output_tokens": 0}, + ) + assert out["priced"] is True + assert _isclose(out["input_usd"], 100_000 / 1_000_000.0 * 5.0) + + +def test_prefix_match_still_resolves_legit_dated_snapshots(): + # Boundary fix must not regress legit dated snapshots. + out = calculate_cost( + "openai", + "gpt-5.4-mini-2026-04-23", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is True + assert _isclose(out["input_usd"], 0.75) + + # And Anthropic dated snapshot still resolves to canonical row. + out = calculate_cost( + "anthropic", + "claude-opus-4-7-20260414", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is True + assert _isclose(out["input_usd"], 5.0) + + +# ── precedence: input_tokens wins over prompt_tokens (and 0 is real) ── + + +def test_explicit_zero_input_tokens_wins_over_stale_prompt_tokens(): + # Input-side mirror of the output zero precedence test. + out = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": 0, + "prompt_tokens": 1_000_000, # stale chat-style mirror + "output_tokens": 100, + }, + ) + assert out["billable_input_tokens"] == 0 + assert out["input_usd"] == 0.0 + + +def test_none_input_tokens_falls_through_to_prompt_tokens(): + # `None` is "key present but unset"; chat-style mirror wins. + out = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": None, + "prompt_tokens": 200_000, + "output_tokens": None, + "completion_tokens": 5_000, + }, + ) + assert out["billable_input_tokens"] == 200_000 + assert out["billable_output_tokens"] == 5_000 + assert _isclose(out["input_usd"], 200_000 / 1_000_000.0 * 5.0) + assert _isclose(out["output_usd"], 5_000 / 1_000_000.0 * 30.0) + + +# ── negative / corrupted upstream values clamp to zero ────────────── + + +def test_negative_tokens_clamp_to_zero_no_negative_bill(): + out = calculate_cost( + "openai", + "gpt-5.5", + {"input_tokens": -100, "output_tokens": -50}, + ) + assert out["billable_input_tokens"] == 0 + assert out["billable_output_tokens"] == 0 + assert out["input_usd"] == 0.0 + assert out["output_usd"] == 0.0 + assert out["total_usd"] == 0.0 + + +def test_negative_cache_buckets_clamp_to_zero(): + # Negative cache_read on Anthropic would otherwise refund the bill. + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 1_000, + "output_tokens": 0, + "cache_creation_input_tokens": -500, + "cache_read_input_tokens": -1_000, + }, + ) + assert out["cache_write_usd"] == 0.0 + assert out["cache_read_usd"] == 0.0 + assert out["billable_input_tokens"] == 1_000 + assert out["total_usd"] >= 0.0 + + +def test_negative_prompt_tokens_chat_style_clamp(): + out = calculate_cost( + "openai", + "gpt-5.4-mini", + {"prompt_tokens": -100, "completion_tokens": -50}, + ) + assert out["billable_input_tokens"] == 0 + assert out["billable_output_tokens"] == 0 + assert out["total_usd"] == 0.0 + + +# ── cache_read > prompt_tokens corruption: no negative billable ───── + + +def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable(): + # cache_read > prompt_tokens clamps uncached_input at 0; billable + # still reflects cache buckets (we charge for what we got). + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "prompt_tokens": 100, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 500, + "completion_tokens": 0, + }, + ) + assert out["input_usd"] == 0.0 # uncached clamped to 0 + assert out["billable_input_tokens"] == 500 # 0 uncached + 500 cache_read + # cache_read still priced at the discount rate. + base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] + assert _isclose( + out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT + ) + + +def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached(): + # OpenAI variant: cached > input must not produce negative input_usd. + base = OPENAI_PRICING["gpt-5.5"]["input_per_mtok"] + out = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": 100, + "output_tokens": 0, + "input_tokens_details": {"cached_tokens": 500}, + }, + ) + assert out["input_usd"] == 0.0 + # Cache read still priced (the 0.1x bucket). + assert _isclose( + out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT + ) + + +# ── long-context tier crosses on billable, including cache_creation ── + + +def test_openai_long_context_triggers_on_cache_creation_inflated_billable(): + # cache_creation pushes billable past 272k -> long-context tier + # must fire to avoid undercounting. + out = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": 250_000, + "cache_creation_input_tokens": 50_000, + "output_tokens": 1_000, + }, + ) + assert out["billable_input_tokens"] == 300_000 + assert "long-context" in out["model_priced"] + assert _isclose(out["input_usd"], 250_000 / 1_000_000.0 * 10.0) + assert _isclose(out["output_usd"], 1_000 / 1_000_000.0 * 45.0) + + +def test_openai_long_context_threshold_boundary_inclusive(): + # Threshold is inclusive (>=). + out = calculate_cost( + "openai", + "gpt-5.5", + {"input_tokens": 272_000, "output_tokens": 1_000}, + ) + assert "long-context" in out["model_priced"] + out_lo = calculate_cost( + "openai", + "gpt-5.5", + {"input_tokens": 271_999, "output_tokens": 1_000}, + ) + assert "long-context" not in out_lo["model_priced"] + + +# ── chat-style vs raw envelope parity at OpenAI long-context tier ── + + +def test_openai_chat_envelope_long_context_parity_with_raw(): + raw = calculate_cost( + "openai", + "gpt-5.5", + {"input_tokens": 300_000, "output_tokens": 10_000}, + ) + chat = calculate_cost( + "openai", + "gpt-5.5", + {"prompt_tokens": 300_000, "completion_tokens": 10_000}, + ) + assert _isclose(chat["total_usd"], raw["total_usd"]) + assert "long-context" in chat["model_priced"] + assert "long-context" in raw["model_priced"] + + +# ── malformed sub-objects: no crash, no false bill ────────────────── + + +def test_cache_creation_as_int_does_not_crash(): + # Proxies sometimes fold cache_creation to an int; tolerate it + # and fall back to the 5m default. + base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"] + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + { + "input_tokens": 0, + "output_tokens": 0, + "cache_creation_input_tokens": 1_000_000, + "cache_creation": 12345, # malformed; must not raise + }, + ) + # Falls back to 5m default for the whole bucket. + assert _isclose( + out["cache_write_usd"], + 1_000_000 / 1_000_000.0 * base * ANTHROPIC_CACHE_5M_WRITE_MULT, + ) + + +def test_non_dict_server_tool_use_is_ignored(): + out = calculate_cost( + "anthropic", + "claude-opus-4-7", + {"input_tokens": 100, "output_tokens": 100, "server_tool_use": "garbage"}, + ) + assert out["server_tools_usd"] == 0.0 + + out = calculate_cost( + "openai", + "gpt-5.5", + {"input_tokens": 100, "output_tokens": 100, "openai_tool_use": [1, 2, 3]}, + ) + assert out["server_tools_usd"] == 0.0 + + +def test_non_dict_input_tokens_details_is_ignored(): + out = calculate_cost( + "openai", + "gpt-5.5", + { + "input_tokens": 100, + "output_tokens": 0, + "input_tokens_details": "nope", + "prompt_tokens_details": [1, 2, 3], + }, + ) + # No cached_tokens recovered -> no discount. + assert out["cache_read_usd"] == 0.0 + + +# ── unknown provider degrades gracefully ──────────────────────────── + + +def test_unknown_provider_priced_false_zero_bill(): + out = calculate_cost( + "gemini", + "gemini-pro", + {"input_tokens": 1_000_000, "output_tokens": 1_000_000}, + ) + assert out["priced"] is False + assert out["total_usd"] == 0.0 + # Tokens still report for the UI. + assert out["billable_input_tokens"] == 1_000_000 + assert out["billable_output_tokens"] == 1_000_000 + + +def test_anthropic_provider_with_openai_model_priced_false(): + # OpenAI id against Anthropic table must not falsely match. + out = calculate_cost( + "anthropic", + "gpt-5.5", + {"input_tokens": 1_000_000, "output_tokens": 0}, + ) + assert out["priced"] is False + + +# ── all-zero / empty usage stays at zero ──────────────────────────── + + +def test_empty_usage_dict_zero_bill(): + out = calculate_cost("openai", "gpt-5.5", {}) + assert out["priced"] is True # model is in the table + assert out["billable_input_tokens"] == 0 + assert out["total_usd"] == 0.0 + + +# ── Defense-in-depth: Anthropic prompt_tokens_details.cached_tokens ── + + +def test_anthropic_prompt_tokens_details_fallback_when_native_key_missing(): + """Chat-style envelope without `cache_read_input_tokens` but with + mirrored `prompt_tokens_details.cached_tokens` should still apply + the cache_read discount.""" + r = calculate_cost( + provider = "anthropic", + model = "claude-opus-4-7", + usage = { + "prompt_tokens": 1_000_000, + "completion_tokens": 0, + # Only the mirrored shape (no native key). + "prompt_tokens_details": {"cached_tokens": 1_000_000}, + "cache_creation_input_tokens": 0, + }, + ) + assert r["billable_input_tokens"] == 1_000_000, r + # 1M cached at 0.1x of $5 (opus 4.7) = $0.50 + assert math.isclose(r["cache_read_usd"], 0.5, rel_tol = 1e-3), r + + +def test_anthropic_native_key_takes_precedence_over_mirrored(): + """When both native and mirrored cache-read fields are present, + the native Anthropic field wins (mirror is fallback-only).""" + r = calculate_cost( + provider = "anthropic", + model = "claude-opus-4-7", + usage = { + "prompt_tokens": 1_000_000, + "cache_read_input_tokens": 800_000, + "prompt_tokens_details": {"cached_tokens": 1_000_000}, + "cache_creation_input_tokens": 0, + }, + ) + # billable = uncached_input + cache_creation + cache_read + # = (1M - 0 - 800k) + 0 + 800k = 1M + assert r["billable_input_tokens"] == 1_000_000, r + # cache_read uses 800k (native), not 1M (mirrored). + assert math.isclose(r["cache_read_usd"], 0.4, rel_tol = 1e-3), r + + +def test_anthropic_native_zero_takes_precedence_over_mirrored(): + """Explicit `cache_read_input_tokens: 0` is authoritative; a stale + mirrored block from a proxy must not inflate cache_read past it.""" + r = calculate_cost( + provider = "anthropic", + model = "claude-opus-4-7", + usage = { + "input_tokens": 1_000_000, + "output_tokens": 0, + "cache_read_input_tokens": 0, + # Stale mirror from a proxy; must be ignored (native present). + "prompt_tokens_details": {"cached_tokens": 1_000_000}, + }, + ) + # Native is 0 -> cache_read stays 0. + assert r["cache_read_usd"] == 0.0, r + # billable = input + cache_creation + cache_read = 1M + 0 + 0 + assert r["billable_input_tokens"] == 1_000_000, r + # 1M uncached at $5/M (no discount). + assert math.isclose(r["input_usd"], 5.0, rel_tol = 1e-3), r + assert math.isclose(r["total_usd"], 5.0, rel_tol = 1e-3), r + + +# ── _build_usage_chunk preserves cache_creation breakdown ── + + +def test_build_usage_chunk_forwards_anthropic_cache_creation_breakdown(): + """Chat-style envelope must carry the 5m/1h cache-write breakdown + so downstream cost calc applies the 2x 1h premium.""" + import json + from core.inference.external_provider import _build_usage_chunk + + chunk = _build_usage_chunk( + completion_id = "cmpl-x", + provider = "anthropic", + last_usage = { + "input_tokens": 10, + "output_tokens": 5, + "cache_creation_input_tokens": 1_000_000, + "cache_read_input_tokens": 0, + "cache_creation": { + "ephemeral_5m_input_tokens": 250_000, + "ephemeral_1h_input_tokens": 750_000, + }, + }, + ) + assert chunk is not None + payload = json.loads(chunk.split("data: ", 1)[1]) + cc = payload["usage"]["cache_creation"] + assert cc["ephemeral_1h_input_tokens"] == 750_000, cc + assert cc["ephemeral_5m_input_tokens"] == 250_000, cc + + +def test_calculate_cost_uses_forwarded_cache_creation_for_1h_premium(): + """Re-emitted chat envelope must price 1h cache writes at 2x base.""" + r = calculate_cost( + provider = "anthropic", + model = "claude-opus-4-7", + usage = { + "prompt_tokens": 1_000_010, + "completion_tokens": 0, + "cache_creation_input_tokens": 1_000_000, + "cache_read_input_tokens": 0, + "cache_creation": { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 1_000_000, + }, + }, + ) + # 1M at 1h-premium (2x of $5 = $10); 5m baseline would be $6.25. + assert math.isclose(r["cache_write_usd"], 10.0, rel_tol = 1e-2), r 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/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py new file mode 100644 index 0000000000..8b90a46d5a --- /dev/null +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -0,0 +1,263 @@ +# 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 `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call +XML that leaks past the speculative buffer in core/inference/llama_cpp.py +when the open/close pair is split across the visible/DRAIN boundary. +""" + +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) + +# Extract the regex from source (routes module needs heavy stubbing to import). +import re as _re + +_src = (Path(_BACKEND_DIR) / "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 source" +_ns = {"_re": _re} +exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) +_TOOL_XML_RE = _ns["_TOOL_XML_RE"] + + +# ── Well-formed pairs ───────────────────────────────────────────── + + +def test_strips_well_formed_tool_call(): + text = ( + "Let me search.\n" + "\n" + "\n" + "\nBillboard 2015\n\n" + "\n" + "\n" + "Here are the songs:" + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "" not in cleaned + assert "" not in cleaned + assert "Here are the songs:" in cleaned, "non-XML content must survive" + assert "Let me search." in cleaned + + +def test_strips_function_only_well_formed(): + text = "Setup.\n\n\nprint(1)\n\n\nDone." + cleaned = _TOOL_XML_RE.sub("", text) + assert "" + "\n" + "\n" + "\nBillboard 2015\n\n" + "" not in cleaned + assert "\n\nprint(1)\n" + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "") + assert "" not in cleaned + assert "Search starting." in cleaned + + +def test_strips_multiple_orphans(): + text = ( + "First call:\n\n\n\nx=1\n" + "Second call:\n\n\nhi\n" + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "" not in cleaned + assert "" not in cleaned + # Mid-string intentionally preserved (see preserve test). + + +# ── Tail-only (PR #5735 follow-up) ─────────────────── + + +def test_strips_tail_only_parameter_orphan(): + # Outer truncated by EOS, inner DRAINED. + cleaned = _TOOL_XML_RE.sub("", "and the text is not readable.\n\n\n") + assert "" not in cleaned + assert "and the text is not readable." in cleaned + + +def test_strips_tail_only_parameter_orphan_single_newline(): + cleaned = _TOOL_XML_RE.sub("", "Global Economic Prospects\n\n") + assert "" not in cleaned + assert "Global Economic Prospects" in cleaned + + +def test_strips_tail_only_parameter_orphan_no_trailing_ws(): + cleaned = _TOOL_XML_RE.sub("", "Final answer.") + assert "" not in cleaned + assert "Final answer." in cleaned + + +def test_preserves_mid_string_parameter_in_code_sample(): + # Tail-anchor on `` is required so doc/example prose survives. + text = ( + "Here is the Qwen tool-call format:\n" + "```xml\n" + "value\n" + "```\n" + "Note the closing sits inside ." + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "Note the closing sits inside" in cleaned + + +def test_strips_well_formed_then_orphan(): + text = ( + "Round one:\n\n\n\n1\n" + "\n\n\n" + "Now round two:\n\n\n\n" + "what is X\n\n" not in cleaned + assert "\n\n\n"Billboard Hot 100" "2015" "weekly" "chart" "position" "3"\n\n\n\n\n"peaked at number 3" Billboard Hot 100 2015 list\n\n\n\n\n"List of Billboard Hot 100 top-ten singles in 2015" wikipedia\n\n\n\nThe user wants me to list and categorize all songs that charted #3 on the Billboard Hot 100 in 2015. I have been trying to get this data", + # Qwen3.6-35B-A3B Q8_0 billboard s21 -- orphan close + "parse it more carefully.\n\n\nThe user wants a list of songs that charted #3 on the Billboard Hot 100 in 2015, categorized.", +] + + +@pytest.mark.parametrize( + "leak", REAL_LEAKS, ids = [f"sweep_sample_{i}" for i in range(len(REAL_LEAKS))] +) +def test_real_world_sweep_leaks_get_stripped(leak): + cleaned = _TOOL_XML_RE.sub("", leak) + assert "" not in cleaned, f"leak survived: {cleaned!r}" + assert " from gdpval sweep ────────── + + +# All end-anchored: outer truncated by EOS, +# inner open DRAINED, leaving bare tail. +GDPVAL_PARAMETER_LEAKS = [ + # Qwen3.5-27B Q8_0 / worldbank s00 + "the page contains image data and the text is not readable.\n\n\n", + # Qwen3.5-27B Q8_0 / worldbank s42 (preceded by mojibake) + "...some mojibake content here...\n\n\n", + # Qwen3.5-27B UD-Q4_K_XL / coppa s07 + "blocked, while others may still be in effect. The law is currently under further review by the Ninth Circuit.\n\n\n", + # Qwen3.5-27B UD-Q4_K_XL / police_training s00 + "comprehensive training report\n\n\n", + # Qwen3.5-27B UD-Q4_K_XL / worldbank s00 + "Global Economic Prospects\nJune 2025\nGlobal Economic Prospects\n\n", + # Qwen3.6-27B Q8_0 / overpass s07 + "Let me create a comprehensive query and instructions document.\n\n\n", +] + + +@pytest.mark.parametrize( + "leak", + GDPVAL_PARAMETER_LEAKS, + ids = [f"gdpval_param_orphan_{i}" for i in range(len(GDPVAL_PARAMETER_LEAKS))], +) +def test_gdpval_parameter_orphans_get_stripped(leak): + cleaned = _TOOL_XML_RE.sub("", leak) + assert "" not in cleaned, f"leak survived: {cleaned!r}" + + +# ── Backtracking guards ────────────────────────────────────────── + + +def test_no_catastrophic_backtracking_on_open_bracket_spam(): + # 256KB of '<' must fail fast (literal mismatch char 2), not backtrack. + import time + + adv = "<" * (1024 * 256) + "X" + t0 = time.perf_counter() + _TOOL_XML_RE.sub("", adv) + elapsed = time.perf_counter() - t0 + assert elapsed < 0.5, f"regex took {elapsed*1000:.0f}ms on 256KB '<' spam" + + +def test_no_catastrophic_backtracking_on_orphan_opening_spam(): + # 1000 unclosed openings: first alt must consume them all greedily. + import time + + adv = "X" * 1000 + t0 = time.perf_counter() + cleaned = _TOOL_XML_RE.sub("", adv) + elapsed = time.perf_counter() - t0 + assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens" + assert "" not in cleaned 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/hardware.py b/studio/backend/utils/hardware/hardware.py index 3764e38272..ede37e2953 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -144,7 +144,10 @@ def detect_hardware() -> DeviceType: if is_apple_silicon() and _has_mlx(): DEVICE = DeviceType.MLX CHAT_ONLY = False - chip = platform.processor() or platform.machine() + # platform.processor() runs `uname -p` which returns "i386" on most + # universal2 / Rosetta-shaped Python builds even on native arm64. + # platform.machine() is "arm64" once is_apple_silicon() has gated us. + chip = platform.machine() or "arm64" print(f"Hardware detected: MLX — Apple Silicon ({chip})") return DEVICE @@ -279,13 +282,11 @@ def get_gpu_memory_info() -> Dict[str, Any]: try: info = mx.device_info() - gpu_name = ( - info.get("device_name") - or platform.processor() - or platform.machine() - ) + # See detect_hardware(): platform.processor() can return "i386" + # on native arm64 Python builds, so prefer machine() as fallback. + gpu_name = info.get("device_name") or platform.machine() or "arm64" except Exception: - gpu_name = platform.processor() or platform.machine() + gpu_name = platform.machine() or "arm64" return { "available": True, diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 4ae9774658..58bc2d8f09 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -14,8 +14,9 @@ "typecheck": "tsc -b --pretty false", "test": "vitest run", "test:watch": "vitest", - "biome:check": "biome check .", - "biome:fix": "biome check . --write" + "i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts", + "biome:check": "biome check", + "biome:fix": "biome check --write" }, "dependencies": { "@assistant-ui/core": "0.1.17", diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index cbd7edf339..4685afa42e 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"; @@ -33,7 +34,9 @@ const routeTree = rootRoute.addChildren([ ]); function DefaultNotFound() { + const t = useT(); const pathname = useRouterState({ select: (s) => s.location.pathname }); + return (

- Page not found + {t("shell.notFound.title")}

- {pathname} does not exist. + {t("shell.notFound.description", { path: pathname })}

); diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 69f9c26fb0..2de179668b 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -7,8 +7,9 @@ import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { IngestionToastStack } from "@/features/rag/components/ingestion-toast-stack"; 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, @@ -17,24 +18,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 = ( -
- Loading... -
-); +function RouteFallback() { + const t = useT(); + + return ( +
+ {t("common.loading")} +
+ ); +} const CHAT_ONLY_ALLOWED = new Set([ "/", @@ -69,6 +71,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"); @@ -76,24 +79,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}` @@ -118,7 +117,7 @@ function RootLayout() { {hideNavbar ? (
- + }>
@@ -144,7 +143,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..849e017ea8 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -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")} > @@ -397,7 +416,7 @@ export function AppSidebar() { sideOffset={6} className="tooltip-compact" > - Close sidebar + {t("shell.aria.closeSidebar")} )} @@ -412,7 +431,7 @@ export function AppSidebar() { type="button" onClick={togglePinned} className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-fg 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="Open sidebar" + aria-label={t("shell.aria.openSidebar")} > @@ -422,7 +441,7 @@ export function AppSidebar() { sideOffset={8} className="tooltip-compact" > - Open sidebar + {t("shell.aria.openSidebar")} @@ -434,7 +453,7 @@ export function AppSidebar() { { @@ -446,7 +465,7 @@ export function AppSidebar() { /> i.id === search.compare)} disabled={chatDisabled} dataTour="chat-compare" @@ -459,7 +478,7 @@ export function AppSidebar() { /> { @@ -477,7 +496,7 @@ export function AppSidebar() { { @@ -489,7 +508,7 @@ export function AppSidebar() { { navigate({ to: "/data-recipes" }); @@ -499,7 +518,7 @@ export function AppSidebar() { { @@ -513,13 +532,16 @@ export function AppSidebar() { - {/* Recent Chats — hide on Studio only (Eyera fac13); chatOpen = ec695 clickability */} {!isStudioRoute && chatItems.length > 0 && ( - + - Recents + {t("shell.navigation.recents")} @@ -552,7 +574,7 @@ export function AppSidebar() { @@ -844,7 +872,9 @@ export function AppSidebar() { - {renamingTarget?.kind === "run" ? "Rename run" : "Rename chat"} + {renamingTarget?.kind === "run" + ? t("shell.dialog.renameRun.title") + : t("shell.dialog.renameChat.title")} @@ -868,14 +906,14 @@ export function AppSidebar() { variant="ghost" onClick={() => setRenamingTarget(null)} > - Cancel + {t("common.cancel")} diff --git a/studio/frontend/src/components/assistant-ui/attachment.tsx b/studio/frontend/src/components/assistant-ui/attachment.tsx index 3ae1c68561..d814490c1d 100644 --- a/studio/frontend/src/components/assistant-ui/attachment.tsx +++ b/studio/frontend/src/components/assistant-ui/attachment.tsx @@ -238,7 +238,7 @@ export const ComposerAddAttachment: FC = () => { side="bottom" variant="ghost" size="icon" - className="aui-composer-add-attachment size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:border-muted-foreground/15 dark:hover:bg-muted-foreground/30" + className="aui-composer-add-attachment size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:hover:bg-muted-foreground/30" aria-label="Add Attachment" > diff --git a/studio/frontend/src/components/assistant-ui/generated-image-overlay-context.tsx b/studio/frontend/src/components/assistant-ui/generated-image-overlay-context.tsx new file mode 100644 index 0000000000..d48c920da1 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/generated-image-overlay-context.tsx @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { + type ReactNode, + createContext, + useCallback, + useContext, + useMemo, + useState, +} from "react"; + +export type GeneratedImageOverlayState = { + image: string; + title: string; + metadata: string; + filename?: string; + openaiImageGenerationCallId?: string; + openaiResponseId?: string; + openaiReasoningItem?: unknown; + threadId?: string | null; +}; + +type GeneratedImageOverlayContextValue = { + overlay: GeneratedImageOverlayState | null; + openOverlay: (overlay: GeneratedImageOverlayState) => void; + closeOverlay: () => void; +}; + +const GeneratedImageOverlayContext = + createContext(null); + +export function GeneratedImageOverlayProvider({ + children, + threadId = null, +}: { + children: ReactNode; + threadId?: string | null; +}) { + const [overlay, setOverlay] = useState( + null, + ); + + const openOverlay = useCallback( + (nextOverlay: GeneratedImageOverlayState) => { + setOverlay({ ...nextOverlay, threadId: nextOverlay.threadId ?? threadId }); + }, + [threadId], + ); + + const closeOverlay = useCallback(() => { + setOverlay(null); + }, []); + + const value = useMemo( + () => ({ overlay, openOverlay, closeOverlay }), + [closeOverlay, openOverlay, overlay], + ); + + return ( + + {children} + + ); +} + +export function useGeneratedImageOverlay(): GeneratedImageOverlayContextValue { + const context = useContext(GeneratedImageOverlayContext); + if (!context) { + throw new Error( + "useGeneratedImageOverlay must be used within GeneratedImageOverlayProvider.", + ); + } + return context; +} diff --git a/studio/frontend/src/components/assistant-ui/image.tsx b/studio/frontend/src/components/assistant-ui/image.tsx new file mode 100644 index 0000000000..a2e3f30dc1 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/image.tsx @@ -0,0 +1,510 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +// +// Portions adapted from assistant-ui packages/ui/src/components/assistant-ui/image.tsx +// MIT License, Copyright (c) 2025 AgentbaseAI Inc. +// Source: https://github.com/assistant-ui/assistant-ui/blob/main/packages/ui/src/components/assistant-ui/image.tsx + +"use client"; + +import { cn } from "@/lib/utils"; +import type { + ImageMessagePart, + ImageMessagePartComponent, +} from "@assistant-ui/react"; +import { type VariantProps, cva } from "class-variance-authority"; +import { + CopyIcon, + DownloadIcon, + ImageIcon, + ImageOffIcon, + Loader2Icon, + RefreshCwIcon, + ShieldAlertIcon, +} from "lucide-react"; +import { + type ComponentProps, + type PropsWithChildren, + memo, + useEffect, + useRef, + useState, +} from "react"; +import { createPortal } from "react-dom"; + +const extensionForMimeType = (mimeType?: string): string => { + switch (mimeType) { + case "image/png": + return "png"; + case "image/jpeg": + case "image/jpg": + return "jpg"; + case "image/webp": + return "webp"; + case "image/gif": + return "gif"; + case "image/svg+xml": + return "svg"; + default: + return "png"; + } +}; + +const DATA_URI_MIME_RE = /data:([^;]+)/; +const DATA_URI_BASE64_RE = /;base64/i; +const IMAGE_DATA_URI_MIME_RE = /^data:([^;,]+)/; + +export const dataUriToBlob = (dataUri: string): Blob => { + const [meta, data] = dataUri.split(","); + const mime = meta?.match(DATA_URI_MIME_RE)?.[1] ?? "application/octet-stream"; + if (!DATA_URI_BASE64_RE.test(meta ?? "")) { + return new Blob([decodeURIComponent(data ?? "")], { type: mime }); + } + const bytes = atob(data ?? ""); + const arr = new Uint8Array(bytes.length); + for (let i = 0; i < bytes.length; i += 1) { + arr[i] = bytes.charCodeAt(i); + } + return new Blob([arr], { type: mime }); +}; + +const mimeFromImage = (image: string): string | undefined => + image.match(IMAGE_DATA_URI_MIME_RE)?.[1]; + +export const downloadImagePart = ( + part: Pick, +): void => { + if (typeof document === "undefined") { + return; + } + const ext = extensionForMimeType(mimeFromImage(part.image)); + const filename = part.filename ?? `image.${ext}`; + const isDataUri = part.image.startsWith("data:"); + const objectUrl = isDataUri + ? URL.createObjectURL(dataUriToBlob(part.image)) + : null; + const href = objectUrl ?? part.image; + const a = document.createElement("a"); + a.href = href; + a.download = filename; + a.rel = "noopener"; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + if (objectUrl) { + URL.revokeObjectURL(objectUrl); + } +}; + +export const copyImagePart = async ( + part: Pick, +): Promise => { + if ( + typeof navigator === "undefined" || + !navigator.clipboard || + typeof ClipboardItem === "undefined" + ) { + throw new Error("Clipboard API is not available in this environment."); + } + const blob = part.image.startsWith("data:") + ? dataUriToBlob(part.image) + : await fetch(part.image).then((r) => r.blob()); + const mime = mimeFromImage(part.image) || blob.type || "image/png"; + await navigator.clipboard.write([new ClipboardItem({ [mime]: blob })]); +}; + +const reportImageCopyError = (error: unknown): void => { + if (typeof window === "undefined") { + return; + } + window.dispatchEvent( + new CustomEvent("assistant-ui:image-copy-error", { detail: error }), + ); +}; + +const imageVariants = cva( + "aui-image-root relative overflow-hidden rounded-lg", + { + variants: { + variant: { + outline: "border border-border", + ghost: "", + muted: "bg-muted/50", + }, + size: { + sm: "max-w-64", + default: "max-w-96", + lg: "max-w-[512px]", + full: "w-full", + }, + }, + defaultVariants: { + variant: "outline", + size: "default", + }, + }, +); + +export type ImageRootProps = ComponentProps<"div"> & + VariantProps; + +function ImageRoot({ + className, + variant, + size, + children, + ...props +}: ImageRootProps) { + return ( +
+ {children} +
+ ); +} + +type ImagePreviewProps = Omit, "children"> & { + containerClassName?: string; +}; + +function ImagePreview({ + className, + containerClassName, + onLoad, + onError, + alt = "Image content", + src, + ...props +}: ImagePreviewProps) { + const imgRef = useRef(null); + const [loadedSrc, setLoadedSrc] = useState(undefined); + const [errorSrc, setErrorSrc] = useState(undefined); + + const loaded = loadedSrc === src; + const error = errorSrc === src; + + useEffect(() => { + if ( + typeof src === "string" && + imgRef.current?.complete && + imgRef.current.naturalWidth > 0 + ) { + setLoadedSrc(src); + } + }, [src]); + + return ( +
+ {!(loaded || error) && ( +
+ +
+ )} + {error ? ( +
+ +
+ ) : ( + {alt} { + if (typeof src === "string") { + setLoadedSrc(src); + } + onLoad?.(e); + }} + onError={(e) => { + if (typeof src === "string") { + setErrorSrc(src); + } + onError?.(e); + }} + /> + )} +
+ ); +} + +function ImageFilename({ + className, + children, + ...props +}: ComponentProps<"span">) { + if (!children) { + return null; + } + + return ( + + {children} + + ); +} + +type ImageZoomProps = PropsWithChildren<{ + src: string; + alt?: string; +}>; + +function ImageZoom({ src, alt = "Image preview", children }: ImageZoomProps) { + const [isOpen, setIsOpen] = useState(false); + + const handleOpen = () => setIsOpen(true); + const handleClose = () => setIsOpen(false); + + useEffect(() => { + if (!isOpen) { + return; + } + const handleKeyDown = (e: KeyboardEvent) => { + if (e.key === "Escape") { + setIsOpen(false); + } + }; + document.addEventListener("keydown", handleKeyDown); + return () => document.removeEventListener("keydown", handleKeyDown); + }, [isOpen]); + + useEffect(() => { + if (!isOpen) { + return; + } + const originalOverflow = document.body.style.overflow; + document.body.style.overflow = "hidden"; + return () => { + document.body.style.overflow = originalOverflow; + }; + }, [isOpen]); + + return ( + <> + + {isOpen && + typeof document !== "undefined" && + createPortal( + , + document.body, + )} + + ); +} + +function ImageGenerating({ className }: { className?: string }) { + return ( +
+ + Generating image… +
+ ); +} + +function ImageContentFilterError({ + className, + reason, +}: { + className?: string; + reason?: string; +}) { + return ( +
+ +

Image could not be generated

+ {reason &&

{reason}

} +
+ ); +} + +export type ImageActionsProps = { + part: ImageMessagePart; + /** + * Wire to your own generation call to show a regenerate button. The button + * renders only when this is set and the part carries a `prompt`. + */ + onRegenerate?: () => void | Promise; + className?: string; +}; + +function RegenerateButton({ + onRegenerate, +}: { + onRegenerate: () => void | Promise; +}) { + const [isRegenerating, setIsRegenerating] = useState(false); + return ( + + ); +} + +function ImageActions({ part, onRegenerate, className }: ImageActionsProps) { + return ( +
+ + + {onRegenerate && } +
+ ); +} + +const ImageImpl: ImageMessagePartComponent = (props) => { + const { image, filename, status } = props; + const alt = filename || "Image content"; + + if (status?.type === "running") { + return ( + + + {filename} + + ); + } + + if (status?.type === "incomplete" && status.reason === "content-filter") { + return ( + + + + ); + } + + return ( + + + + + {filename} + + ); +}; + +const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & { + Root: typeof ImageRoot; + Preview: typeof ImagePreview; + Filename: typeof ImageFilename; + Zoom: typeof ImageZoom; + Actions: typeof ImageActions; + Generating: typeof ImageGenerating; + ContentFilterError: typeof ImageContentFilterError; +}; + +Image.displayName = "Image"; +Image.Root = ImageRoot; +Image.Preview = ImagePreview; +Image.Filename = ImageFilename; +Image.Zoom = ImageZoom; +Image.Actions = ImageActions; +Image.Generating = ImageGenerating; +Image.ContentFilterError = ImageContentFilterError; + +export { + Image, + ImageRoot, + ImagePreview, + ImageFilename, + ImageZoom, + ImageActions, + ImageGenerating, + ImageContentFilterError, + imageVariants, +}; diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 4fb68d4b90..31f742bc5e 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -33,10 +33,24 @@ export const MessageTiming: FC<{ if (timing?.totalStreamTime === undefined) return null; - const serverTimings = ( + const custom = ( message.metadata as Record | undefined - )?.custom as { serverTimings?: Record } | undefined; - const st = serverTimings?.serverTimings; + )?.custom as + | { + serverTimings?: Record; + contextUsage?: { + cachedTokens?: number; + cacheWriteTokens?: number; + }; + } + | undefined; + const st = custom?.serverTimings; + // `??` (not `||`) so an explicit cache_n=0 isn't replaced by a stale + // contextUsage.cachedTokens from a prior turn. + const cacheHits = + st?.cache_n ?? custom?.contextUsage?.cachedTokens ?? 0; + // Anthropic-only cache-write count. + const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0; // Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op // turns, blowing the rate up to Infinity. Require >=1 token AND a @@ -122,11 +136,19 @@ export const MessageTiming: FC<{ )} - {(st?.cache_n ?? 0) > 0 && ( + {cacheHits > 0 && (
Cache hits - {formatNumber(st!.cache_n)} + {formatNumber(cacheHits)} + +
+ )} + {cacheWrites > 0 && ( +
+ Cache writes + + {formatNumber(cacheWrites)}
)} @@ -146,7 +168,7 @@ export const MessageTiming: FC<{ ) : ( <> - {/* Client-side metrics (safetensors fallback) */} + {/* Client-side metrics (safetensors + external provider fallback) */} {timing.firstTokenTime !== undefined && (
First token @@ -155,6 +177,22 @@ export const MessageTiming: FC<{
)} + {cacheHits > 0 && ( +
+ Cache hits + + {formatNumber(cacheHits)} + +
+ )} + {cacheWrites > 0 && ( +
+ Cache writes + + {formatNumber(cacheWrites)} + +
+ )}
Total diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index e4401cc12e..285b87d138 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -134,7 +134,7 @@ function ReasoningTrigger({ {active ? ( Thinking... ) : ( - Thought for {duration ?? 0} seconds + Thought for {duration ?? 0} {duration === 1 ? "second" : "seconds"} )} { "url" in part && part.url ) { + const url = part.url as string; + const partId = + typeof (part as { id?: unknown }).id === "string" + ? ((part as { id: string }).id) + : url; sources.push({ kind: "url", - url: part.url as string, + id: partId, + url, title: (part as { title?: string }).title || "", description: (part as { metadata?: { description?: string } }) .metadata?.description, diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index f3c071d8f0..d3cbb2d0a8 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -7,6 +7,11 @@ import { UserMessageAttachments, } from "@/components/assistant-ui/attachment"; import { CodeToggleIcon } from "@/components/assistant-ui/code-toggle-icon"; +import { + GeneratedImageOverlayProvider, + useGeneratedImageOverlay, +} from "@/components/assistant-ui/generated-image-overlay-context"; +import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; @@ -46,13 +51,14 @@ import { import { sentAudioNames } from "@/features/chat/api/chat-adapter"; import { parseExternalModelId } from "@/features/chat/external-providers"; import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities"; -import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; +import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; +import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params"; import { isTauri } from "@/lib/api-base"; -import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message"; import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -88,30 +94,31 @@ import { TerminalIcon, XIcon, } from "lucide-react"; -import { Copy01Icon, Delete02Icon, Edit03Icon, Tick02Icon } from "@hugeicons/core-free-icons"; +import { + Copy01Icon, + Delete02Icon, + Edit03Icon, + Image03Icon, + Tick02Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ChangeEvent, + type ComponentProps, type CompositionEvent, type FC, - type FormEvent, type KeyboardEvent, useCallback, useEffect, useRef, useState, } from "react"; -import { toast } from "@/lib/toast"; export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean; targetThreadId?: string; -}> = ({ - hideComposer, - hideWelcome, - targetThreadId, -}) => { +}> = ({ hideComposer, hideWelcome, targetThreadId }) => { // Intent-aware autoscroll: replaces assistant-ui's built-in autoscroll // to prevent the streaming-mutation race that makes the viewport snap // back to the bottom while the user is scrolling up (see the hook for @@ -122,85 +129,205 @@ export const Thread: FC<{ const isComposerAttachPending = useAuiState(({ threads }) => targetThreadId ? threads.mainThreadId !== targetThreadId : false, ); + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const threadId = targetThreadId ?? activeThreadId ?? null; return ( - - - - {!hideWelcome && ( - thread.isEmpty && !thread.isLoading}> - - - )} + + + + + {!hideWelcome && ( + thread.isEmpty && !thread.isLoading} + > + + + )} - + - {/* Bottom slack so the last message has breathing room above the + {/* Bottom slack so the last message has breathing room above the sticky scroll-to-bottom button (and the floating composer in single mode). Without this, content would butt against the sticky footer and feel cramped. */} - hideWelcome || !thread.isEmpty}> -
- - - hideWelcome || !thread.isEmpty}> - - - - - - - {!hideComposer && ( - hideWelcome || !thread.isEmpty}> -
+ hideWelcome || !thread.isEmpty}>
-
-
- -
-

- LLMs can make mistakes. Double-check responses. -

-
-
- + + + hideWelcome || !thread.isEmpty}> + + + + + + + + + {!hideComposer && ( + hideWelcome || !thread.isEmpty}> + + + )} + + + + ); +}; + +const GeneratedImageViewportOverlay: FC<{ hideComposer?: boolean }> = ({ + hideComposer, +}) => { + const { overlay, closeOverlay } = useGeneratedImageOverlay(); + + useEffect(() => { + if (!overlay) { + return; + } + document.querySelector(".aui-composer-input")?.focus(); + }, [overlay]); + + if (!overlay) { + return null; + } + + return ( +
+ + +
+
+
+ {overlay.title} +
+
+

+ Generated image +

+ {overlay.metadata ? ( +

+ {overlay.metadata} +

+ ) : null} + {hideComposer ? null : ( +

+ Type edits below, then send +

+ )} +
+
+ +
+ ); +}; + +const ThreadComposerDock: FC<{ + disabled?: boolean; + threadId?: string | null; +}> = ({ disabled, threadId }) => { + const { overlay } = useGeneratedImageOverlay(); + + return ( +
+
+
+
+ +
+

+ LLMs can make mistakes. Double-check responses. +

+
+
); }; @@ -228,13 +355,17 @@ const ThreadScrollToBottom: FC = () => { ); }; -const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { +const ThreadWelcome: FC<{ + hideComposer?: boolean; + threadId?: string | null; +}> = ({ hideComposer, threadId }) => { const [currentEmoji, setCurrentEmoji] = useState("large sloth drink.png"); useEffect(() => { const hour = new Date().getHours(); if (hour >= 6 && hour < 12) setCurrentEmoji("large sloth drink.png"); - else if (hour >= 12 && hour < 17) setCurrentEmoji("sloth magnify final.png"); + else if (hour >= 12 && hour < 17) + setCurrentEmoji("sloth magnify final.png"); else if (hour >= 17 && hour < 21) setCurrentEmoji("sloth shy large.png"); else setCurrentEmoji("unsloth-gem.png"); }, []); @@ -249,11 +380,7 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
- Sloth mascot + Sloth mascot

Chat with your model

@@ -261,18 +388,21 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => { Run GGUFs, safetensors, vision and audio models

- {!hideComposer && } + {!hideComposer && }
); }; -const ComposerAnimated: FC<{ disabled?: boolean }> = ({ disabled }) => { +const ComposerAnimated: FC<{ + disabled?: boolean; + threadId?: string | null; +}> = ({ disabled, threadId }) => { return (
- +
); @@ -302,8 +432,21 @@ const PendingAudioChip: FC = () => { ); }; -const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { - const { inputProps, isComposing, isComposingRef } = useImeComposerInputHandlers(); +const Composer: FC<{ + disabled?: boolean; + threadId?: string | null; +}> = ({ disabled, threadId }) => { + const aui = useAui(); + const { overlay, closeOverlay } = useGeneratedImageOverlay(); + const setImageToolsEnabled = useChatRuntimeStore( + (s) => s.setImageToolsEnabled, + ); + const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); + const setPendingImageEditReference = useChatRuntimeStore( + (s) => s.setPendingImageEditReference, + ); + const { inputProps, isComposing, isComposingRef } = + useImeComposerInputHandlers(); const composerText = useAuiState(({ composer }) => composer.text); const hasAttachments = useAuiState( ({ composer }) => composer.attachments.length > 0, @@ -313,12 +456,23 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { (attachment) => attachment.status.type === "running", ), ); - const hasPendingAudio = useChatRuntimeStore((s) => Boolean(s.pendingAudioName)); + const hasPendingAudio = useChatRuntimeStore((s) => + Boolean(s.pendingAudioName), + ); const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled); const { pendingDocs, addDoc, removeDoc, clearDocs, isIndexing } = useThreadDocUploads(); + const referenceThreadId = threadId ?? activeThreadId ?? null; const hasSendableContent = composerText.trim().length > 0 || hasAttachments || hasPendingAudio; + const shouldBlockSend = useCallback( + () => + !hasSendableContent || + isComposingRef.current || + hasPendingAttachments || + isIndexing, + [hasPendingAttachments, hasSendableContent, isComposingRef, isIndexing], + ); const sendBlocked = disabled || !hasSendableContent || @@ -327,27 +481,67 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { isIndexing; const handleSubmit = useCallback( - (event: FormEvent) => { - if ( - disabled || - !hasSendableContent || - isComposingRef.current || - hasPendingAttachments || - isIndexing - ) { + (event: Parameters["onSubmit"]>>[0]) => { + if (disabled || shouldBlockSend()) { event.preventDefault(); return; } + + if (overlay) { + const trimmed = composerText.trim(); + if (!trimmed) { + event.preventDefault(); + return; + } + if (!overlay.openaiImageGenerationCallId) { + event.preventDefault(); + toast.error("This generated image cannot be edited", { + description: + "The original image reference is missing. Generate the image again, then retry the edit.", + }); + closeOverlay(); + return; + } + if ((overlay.threadId ?? null) !== referenceThreadId) { + event.preventDefault(); + toast.error("This generated image belongs to another chat", { + description: "Open the original chat and retry the edit.", + }); + closeOverlay(); + return; + } + setImageToolsEnabled(true); + setPendingImageEditReference({ + threadId: overlay.threadId ?? referenceThreadId, + openaiImageGenerationCallId: overlay.openaiImageGenerationCallId, + ...(overlay.openaiResponseId + ? { openaiResponseId: overlay.openaiResponseId } + : {}), + openaiReasoningItem: overlay.openaiReasoningItem, + }); + flushResourcesSync(() => { + aui + .composer() + .setText( + `Use the selected generated image as the reference and apply this edit: ${trimmed}. Preserve everything else exactly.`, + ); + }); + closeOverlay(); + } // Drop chips on send; docs stay searchable in the backend. clearDocs(); }, [ - disabled, - hasPendingAttachments, - hasSendableContent, - isComposingRef, - isIndexing, + aui, clearDocs, + closeOverlay, + composerText, + disabled, + overlay, + referenceThreadId, + setImageToolsEnabled, + setPendingImageEditReference, + shouldBlockSend, ], ); @@ -358,13 +552,15 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { = ({ disabled }) => { /> - !hasSendableContent || - isComposingRef.current || - hasPendingAttachments || - isIndexing - } + shouldBlockSend={shouldBlockSend} ragModeOn={ragToolEnabled} onAddDoc={addDoc} /> @@ -586,7 +777,6 @@ const ComposerAudioUpload: FC = () => { ); }; - const ReasoningToggle: FC = () => { const modelLoaded = useChatRuntimeStore( (s) => !!s.params.checkpoint && !s.modelLoading, @@ -598,8 +788,12 @@ const ReasoningToggle: FC = () => { const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled); const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle); const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort); - const supportsReasoningOff = useChatRuntimeStore((s) => s.supportsReasoningOff); - const reasoningEffortLevels = useChatRuntimeStore((s) => s.reasoningEffortLevels); + const supportsReasoningOff = useChatRuntimeStore( + (s) => s.supportsReasoningOff, + ); + const reasoningEffortLevels = useChatRuntimeStore( + (s) => s.reasoningEffortLevels, + ); const setReasoningEffort = useChatRuntimeStore((s) => s.setReasoningEffort); const lastOpenRouterChosenModel = useChatRuntimeStore( (s) => s.lastOpenRouterChosenModel, @@ -631,6 +825,7 @@ const ReasoningToggle: FC = () => { { isReasoningProvider: selectedExternalProvider?.isReasoningModel === true, + baseUrl: selectedExternalProvider?.baseUrl ?? null, }, ) : null; @@ -652,7 +847,8 @@ const ReasoningToggle: FC = () => { effectiveReasoningEnabled && reasoningEffort !== "none"; const disabled = !(modelLoaded && effectiveSupportsReasoning); const formatEffortLabel = (level: typeof reasoningEffort): string => { - if (level !== "xhigh") return level.charAt(0).toUpperCase() + level.slice(1); + if (level !== "xhigh") + return level.charAt(0).toUpperCase() + level.slice(1); const normalized = externalSelection?.modelId?.trim().toLowerCase() ?? ""; if ( normalized.startsWith("claude-opus-4-6") || @@ -710,23 +906,25 @@ const ReasoningToggle: FC = () => { {effectiveReasoningEffortLevels .filter((level) => level !== "none") .map((level) => ( - { - setReasoningEffort(level); - setReasoningEnabled(true); - applyQwenThinkingParams(true); - // Kimi's $web_search builtin forbids thinking, so - // enabling thinking flips the Search pill off. - if (isKimiExternal && toolsEnabled) { - setToolsEnabled(false); - } - }} - > - {formatEffortLabel(level)} - {effectiveReasoningVisualEnabled && reasoningEffort === level ? " \u2713" : ""} - - ))} + { + setReasoningEffort(level); + setReasoningEnabled(true); + applyQwenThinkingParams(true); + // Kimi's $web_search builtin forbids thinking, so + // enabling thinking flips the Search pill off. + if (isKimiExternal && toolsEnabled) { + setToolsEnabled(false); + } + }} + > + {formatEffortLabel(level)} + {effectiveReasoningVisualEnabled && reasoningEffort === level + ? " \u2713" + : ""} + + ))} ); @@ -841,8 +1039,7 @@ const WebSearchToggle: FC = () => { ? externalProviders.find((p) => p.id === externalSelection.providerId) : undefined; const isKimiExternal = selectedExternalProvider?.providerType === "kimi"; - const disabled = - !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); + const disabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); return ( ); @@ -1073,10 +1272,10 @@ const RagDocAttachment: FC<{ onSelect: (file: File) => void }> = ({ const ComposerAction: FC<{ disabled?: boolean; - blockSend?: () => boolean; + shouldBlockSend?: () => boolean; ragModeOn?: boolean; onAddDoc?: (file: File) => void; -}> = ({ disabled, blockSend, ragModeOn, onAddDoc }) => { +}> = ({ disabled, shouldBlockSend, ragModeOn, onAddDoc }) => { return (
@@ -1128,7 +1327,7 @@ const ComposerAction: FC<{ size="icon" disabled={disabled} onClick={(event) => { - if (blockSend?.()) { + if (shouldBlockSend?.()) { event.preventDefault(); } }} @@ -1397,7 +1596,11 @@ const UserActionBar: FC = () => { - + diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index f0d6262a45..2ea87d63f3 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -101,6 +101,16 @@ const statusIconMap: Record = { "requires-action": AlertCircleIcon, }; +const MCP_TOOL_PREFIX = "mcp__"; + +function formatToolNameForDisplay(toolName: string): string { + if (!toolName.startsWith(MCP_TOOL_PREFIX)) return toolName; + const rest = toolName.slice(MCP_TOOL_PREFIX.length); + const sep = rest.indexOf("__"); + if (sep <= 0) return toolName; + return `${rest.slice(0, sep)} · ${rest.slice(sep + 2)}`; +} + function ToolFallbackTrigger({ toolName, status, @@ -119,6 +129,7 @@ function ToolFallbackTrigger({ const StatusIcon = statusIconMap[statusType]; const label = isCancelled ? "Cancelled tool" : "Used tool"; + const displayName = formatToolNameForDisplay(toolName); return ( {label}:{" "} - {toolName} + {displayName} {isRunning && ( {label}:{" "} - {toolName} + {displayName} )} diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx index 246d8b6978..7dfdd903fe 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx @@ -3,9 +3,14 @@ "use client"; -import { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react"; -import { ImageIcon, LoaderIcon } from "lucide-react"; -import { memo, useEffect, useState } from "react"; +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; +import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; +import { DownloadIcon, ImageIcon, PencilIcon } from "lucide-react"; +import type { CSSProperties, MouseEvent } from "react"; +import { memo, useCallback, useEffect, useRef, useState } from "react"; +import { useGeneratedImageOverlay } from "./generated-image-overlay-context"; +import { Image, downloadImagePart } from "./image"; import { ToolFallbackContent, ToolFallbackRoot, @@ -38,6 +43,9 @@ import { interface ImageGenerationArgs { prompt?: string; kind?: string; + openai_image_generation_call_id?: unknown; + openai_response_id?: unknown; + openai_reasoning_item?: unknown; } interface ImageGenerationResult { @@ -46,6 +54,85 @@ interface ImageGenerationResult { size?: string; quality?: string; background?: string; + prompt?: string; +} + +type GeneratedImagePart = { + type: "image"; + image: string; + filename?: string; +}; + +const CAPTION_COLLAPSED_LINES = 4; + +const extensionForMime = (mime: string): string => { + switch (mime.toLowerCase()) { + case "image/jpeg": + case "image/jpg": + return "jpg"; + case "image/webp": + return "webp"; + case "image/gif": + return "gif"; + case "image/svg+xml": + return "svg"; + default: + return "png"; + } +}; + +const imageFilenameFromPrompt = (prompt: string, mime: string): string => { + const slug = prompt + .trim() + .toLowerCase() + .replace(/[^a-z0-9]+/g, "-") + .replace(/^-|-$/g, "") + .slice(0, 48); + return `${slug || "generated-image"}.${extensionForMime(mime)}`; +}; + +const formatGeneratedImageLabel = (prompt: string): string => { + if (!prompt) { + return "Generated image"; + } + return prompt.length > 80 + ? `Generated image: ${prompt.slice(0, 80)}…` + : `Generated image: ${prompt}`; +}; + +const loadingDots = Array.from({ length: 64 }, (_, index) => { + const row = Math.floor(index / 8); + const col = index % 8; + return ( + + ); +}); + +function GeneratedImagePlaceholder({ label }: { label: string }) { + return ( +
+ {label} +
+ {loadingDots} +
+
+ ); } const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ @@ -53,6 +140,7 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ result, status, }) => { + const { openOverlay } = useGeneratedImageOverlay(); const parsedArgs = (args as ImageGenerationArgs) ?? {}; const prompt = parsedArgs.prompt ?? ""; const isRunning = status?.type === "running"; @@ -66,33 +154,131 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ const imageSrc = imageResult?.image_b64 ? `data:${mime};base64,${imageResult.image_b64}` : null; + const imageTitle = + imageResult?.prompt?.trim() || prompt.trim() || "Generated image"; + const captionPrompt = imageResult?.prompt?.trim() || prompt.trim(); + const promptLikelyNeedsExpansion = captionPrompt.length > 220; + const imageMetadata = [imageResult?.size, imageResult?.quality, mime] + .filter(Boolean) + .join(" · "); + const openaiImageGenerationCallId = + typeof parsedArgs.openai_image_generation_call_id === "string" + ? parsedArgs.openai_image_generation_call_id + : undefined; + const openaiResponseId = + typeof parsedArgs.openai_response_id === "string" + ? parsedArgs.openai_response_id + : undefined; + const imagePart: GeneratedImagePart | null = imageSrc + ? { + type: "image", + image: imageSrc, + filename: imageFilenameFromPrompt(prompt, mime), + } + : null; - // Collapse the card once the model has resumed streaming prose - // after the image. Mirrors CodeExecutionToolUI so the inline image - // doesn't collapse mid-stream and the user can click to re-expand. - const hasText = useAuiState(({ message }) => - message.content.some( - (p) => - p.type === "text" && - "text" in p && - (p as { text: string }).text.length > 0, - ), - ); const [open, setOpen] = useState(true); - useEffect(() => { - if (isRunning) { - setOpen(true); - } else if (hasText && !imageSrc) { - setOpen(false); + const [expandedCaptionPrompt, setExpandedCaptionPrompt] = useState< + string | null + >(null); + const [promptOverflow, setPromptOverflow] = useState<{ + prompt: string; + canExpand: boolean; + } | null>(null); + const captionRef = useRef(null); + const isPendingImage = !imagePart && status?.type === "running"; + + const promptOverflowMeasured = promptOverflow?.prompt === captionPrompt; + const promptCanExpand = promptOverflowMeasured + ? promptOverflow.canExpand + : false; + const promptExpanded = expandedCaptionPrompt === captionPrompt; + + const updatePromptOverflow = useCallback(() => { + const captionElement = captionRef.current; + if (!captionElement || !captionPrompt) { + return; } - }, [isRunning, hasText, imageSrc]); + const computedStyle = window.getComputedStyle(captionElement); + const lineHeight = Number.parseFloat(computedStyle.lineHeight); + const collapsedHeight = + (Number.isFinite(lineHeight) ? lineHeight : 20) * + CAPTION_COLLAPSED_LINES; + const hasOverflow = captionElement.scrollHeight > collapsedHeight + 1; + setPromptOverflow((current) => + current?.prompt === captionPrompt && current.canExpand === hasOverflow + ? current + : { prompt: captionPrompt, canExpand: hasOverflow }, + ); + }, [captionPrompt]); + + useEffect(() => { + const captionElement = captionRef.current; + if (!captionElement || !captionPrompt) { + return; + } + const frame = window.requestAnimationFrame(updatePromptOverflow); + const resizeObserver = + typeof ResizeObserver === "undefined" + ? null + : new ResizeObserver(updatePromptOverflow); + resizeObserver?.observe(captionElement); + window.addEventListener("resize", updatePromptOverflow); + return () => { + window.cancelAnimationFrame(frame); + resizeObserver?.disconnect(); + window.removeEventListener("resize", updatePromptOverflow); + }; + }, [captionPrompt, updatePromptOverflow]); + + const shouldClampPrompt = + (promptOverflowMeasured ? promptCanExpand : promptLikelyNeedsExpansion) && + !promptExpanded; const runningLabel = "Generating image…"; - const completedLabel = prompt - ? prompt.length > 80 - ? `Generated image: ${prompt.slice(0, 80)}…` - : `Generated image: ${prompt}` - : "Generated image"; + const completedLabel = formatGeneratedImageLabel(prompt); + + const showPreview = () => { + if (!imagePart) { + return; + } + openOverlay({ + image: imagePart.image, + title: imageTitle, + metadata: imageMetadata, + filename: imagePart.filename, + openaiImageGenerationCallId, + openaiResponseId, + openaiReasoningItem: parsedArgs.openai_reasoning_item, + }); + }; + + const stopOverlayActionPropagation = ( + event: MouseEvent, + ) => { + event.preventDefault(); + event.stopPropagation(); + }; + + const handleDownload = (event: MouseEvent) => { + stopOverlayActionPropagation(event); + if (imagePart) { + downloadImagePart(imagePart); + } + }; + + const handleEditClick = (event: MouseEvent) => { + stopOverlayActionPropagation(event); + showPreview(); + }; + + if (isPendingImage) { + return ( +
+ +
+ ); + } return ( @@ -102,21 +288,77 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ icon={ImageIcon} /> - {isRunning && !imageSrc ? ( -
- - {runningLabel} -
- ) : imageSrc ? ( -
- {prompt - {prompt ? ( -
- {prompt} + {imagePart ? ( +
+
+ +
+ +
+ + +
+
+ {captionPrompt ? ( +
+
+ {captionPrompt} +
+ {promptCanExpand ? ( + + ) : null}
) : null}
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx index da4c0cbbb9..b9bda2e832 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx @@ -24,6 +24,22 @@ const RE_TITLE = /Title:\s*(.+)/; const RE_URL = /URL:\s*(.+)/; const RE_SNIPPET = /Snippet:\s*(.+)/s; +/** + * Reject anything that is not a real http(s) URL. Web-search / web-fetch + * output is provider-controlled, so hostile ``javascript:`` / ``data:`` + * lines must not reach the Source badge's . + */ +function isSafeHttpUrl(raw: string): boolean { + const value = raw.trim(); + if (!value || /[\r\n]/.test(value)) return false; + try { + const parsed = new URL(value); + return parsed.protocol === "http:" || parsed.protocol === "https:"; + } catch { + return false; + } +} + /** Parse the backend's "Title: ...\nURL: ...\nSnippet: ...\n---" format into structured sources. */ function parseSearchResults(raw: string): ParsedSource[] { if (!raw) { @@ -35,13 +51,14 @@ function parseSearchResults(raw: string): ParsedSource[] { const titleMatch = block.match(RE_TITLE); const urlMatch = block.match(RE_URL); const snippetMatch = block.match(RE_SNIPPET); - if (titleMatch && urlMatch) { - sources.push({ - title: titleMatch[1].trim(), - url: urlMatch[1].trim(), - snippet: snippetMatch?.[1]?.trim() ?? "", - }); - } + if (!titleMatch || !urlMatch) continue; + const url = urlMatch[1].trim(); + if (!isSafeHttpUrl(url)) continue; + sources.push({ + title: titleMatch[1].trim(), + url, + snippet: snippetMatch?.[1]?.trim() ?? "", + }); } return sources; } diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index e9fe1d679c..c0a60d0de2 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -108,6 +108,7 @@ export const LR_DEFAULT_CPT = 5e-5; export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, + visionImageSize: null as number | null, learningRate: LR_DEFAULT_LORA, // null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT. embeddingLearningRate: null as number | null, diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 135554d642..2ef884d6e0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -5,7 +5,7 @@ import { type ParsedChunk, parseChunks, } from "@/components/assistant-ui/tool-ui-search-knowledge-base"; -import { getAuthToken } from "@/features/auth/session"; +import { getAuthToken } from "@/features/auth"; import { type SearchHit, type SearchRequest, @@ -32,20 +32,28 @@ import { pickFriendlyContainerName } from "../lib/friendly-names"; import { EXTERNAL_MAX_OUTPUT_TOKENS, clampReasoningEffortToLevels, + getExternalMaxOutputTokens, getExternalMinOutputTokens, getExternalReasoningCapabilities, getProviderCapabilities, + isGeminiCustomOpenAICompatBase, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, providerSupportsBuiltinWebFetch, providerSupportsBuiltinWebSearch, + providerSupportsFastMode, } from "../provider-capabilities"; -import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import { + type PendingImageEditReference, + useChatRuntimeStore, +} from "../stores/chat-runtime-store"; import { useExternalProvidersStore } from "../stores/external-providers-store"; import { isMultimodalResponse } from "../types/api"; import type { OpenAIChatCompletionsRequest, + OpenAIChatMessage, OpenAIMessageContent, + OpenAIReasoningContentPart, } from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; import { @@ -134,6 +142,13 @@ interface ServerUsage { prompt_tokens: number; completion_tokens: number; total_tokens: number; + // External prompt-cache fields (see _build_usage_chunk in + // external_provider.py). cache_creation is Anthropic-only. + prompt_tokens_details?: { + cached_tokens?: number; + }; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; } /** Server-side timing data from llama-server's timings object. */ @@ -155,6 +170,33 @@ type RunMessage = RunMessages[number]; /** Tracks which user messages were sent with an audio file (messageId → filename). */ export const sentAudioNames = new Map(); +// Synthetic provider-side tool names; backend stamps args._server_tool +// so user functions with the same name aren't dropped. Mirror of +// _SERVER_SIDE_BUILTIN_TOOL_NAMES on the backend. +const SERVER_SIDE_BUILTIN_TOOL_NAMES = new Set([ + "web_search", + "web_fetch", + "code_execution", + "image_generation", +]); + +/** + * Whether a persisted tool-call part is provider-side synthetic and + * should be stripped from outbound history. Match on the + * args._server_tool marker or a Gemini native_part payload — no shape + * heuristic, because user functions can legitimately share a name. + */ +function isServerSideBuiltinToolPart( + toolNameLower: string, + _argsObj: Record | null, + hasServerToolMarker: boolean, + hasNativePart: boolean, +): boolean { + if (!SERVER_SIDE_BUILTIN_TOOL_NAMES.has(toolNameLower)) return false; + if (hasServerToolMarker) return true; + return hasNativePart; +} + /** * Match error messages that indicate the request filled or would fill * the KV cache, so the UI can show a dedicated toast pointing at the @@ -207,6 +249,91 @@ async function updateStoredChatThreadEventually( } } +/** + * Return ``raw`` when it is a safe-to-navigate http(s) URL, or "" otherwise. + * Rejects non-string input, CR/LF (header injection), and non-http(s) + * schemes (``javascript:`` / ``data:`` / ``vbscript:``) so provider / + * tool-controlled strings cannot land in an . + */ +function isSafeNavigableSourceUrl(raw: unknown): string { + if (typeof raw !== "string") return ""; + const value = raw.trim(); + if (!value || /[\r\n]/.test(value)) return ""; + try { + const parsed = new URL(value); + if (parsed.protocol === "http:" || parsed.protocol === "https:") { + return value; + } + } catch { + // Fall through. + } + return ""; +} + +/** Convert an Anthropic document citation dict into a Sources-panel source. */ +function documentCitationToSource( + cit: Record, + fallbackIdx: number, +): { + type: "source"; + sourceType: "url"; + id: string; + url: string; + title: string; + metadata?: { description: string }; +} | null { + const source = + typeof cit.source === "string" && cit.source ? cit.source : ""; + const docTitle = + (typeof cit.document_title === "string" && cit.document_title) || + (typeof cit.title === "string" && cit.title) || + ""; + const docIndex = + typeof cit.document_index === "number" ? cit.document_index : undefined; + // Only treat ``source`` as a navigable URL when it is real http(s); + // search_result_location can carry a free-form id (e.g. ``kb-doc-42``) + // or a hostile ``javascript:`` / ``data:`` / ``vbscript:`` string. + // Fall back to a stable doc anchor otherwise. + const url = + isSafeNavigableSourceUrl(source) || `#anthropic-doc-${docIndex ?? fallbackIdx}`; + const title = docTitle || source || `Document ${fallbackIdx + 1}`; + const cited = + typeof cit.cited_text === "string" ? cit.cited_text.trim() : ""; + // Trim the cited snippet so the Sources panel stays scannable. + const description = + cited.length > 240 ? `${cited.slice(0, 240)}...` : cited; + // Anthropic numbers inline [N] per citation, not per source URL. + // Fold citation type + position-bearing fields into the id so two + // distinct citations on the same source (or two search_result_locations + // with different search_result_index) keep separate Sources entries. + const citationType = + typeof cit.type === "string" ? String(cit.type) : ""; + const positionParts = [ + cit.search_result_index, + cit.start_char_index, + cit.end_char_index, + cit.start_page_number, + cit.end_page_number, + cit.start_block_index, + cit.end_block_index, + ] + .filter((v) => typeof v === "number") + .map((v) => String(v)) + .join(":"); + const idAnchor = positionParts + ? `${citationType}:${positionParts}` + : `${citationType}:${fallbackIdx}`; + const id = `${url}#${idAnchor}`; + return { + type: "source" as const, + sourceType: "url" as const, + id, + url, + title, + ...(description ? { metadata: { description } } : {}), + }; +} + /** Parse "Title: ...\nURL: ...\nSnippet: ..." blocks into source content parts. */ function parseSourcesFromResult(raw: string): { type: "source"; @@ -231,7 +358,11 @@ function parseSourcesFromResult(raw: string): { const urlMatch = block.match(/URL:\s*(.+)/); const snippetMatch = block.match(/Snippet:\s*(.+)/); if (titleMatch && urlMatch) { - const url = urlMatch[1].trim(); + // Drop blocks whose ``URL:`` is not safe http(s); provider/tool + // output is attacker-controllable so a hostile ``javascript:`` / + // ``data:`` line must not reach the Sources panel . + const url = isSafeNavigableSourceUrl(urlMatch[1]); + if (!url) continue; const snippet = snippetMatch?.[1]?.trim(); sources.push({ type: "source" as const, @@ -501,37 +632,30 @@ function collectImageParts( message: RunMessage, ): Array<{ type: "image_url"; image_url: { url: string } }> { const parts: Array<{ type: "image_url"; image_url: { url: string } }> = []; + const pushImagePart = (part: { type: string }) => { + if (part.type !== "image" || !("image" in part)) { + return; + } + const src = (part as { image: string }).image; + if (!src) { + return; + } + parts.push({ + type: "image_url", + image_url: { + url: src.startsWith("data:") ? src : `data:image/png;base64,${src}`, + }, + }); + }; for (const part of message.content ?? []) { - if (part.type === "image" && "image" in part) { - const src = (part as { image: string }).image; - if (src) { - parts.push({ - type: "image_url", - image_url: { - url: src.startsWith("data:") ? src : `data:image/png;base64,${src}`, - }, - }); - } - } + pushImagePart(part); } if ("attachments" in message && (message.attachments?.length ?? 0) > 0) { for (const attachment of message.attachments ?? []) { for (const part of attachment.content ?? []) { - if (part.type === "image" && "image" in part) { - const src = (part as { image: string }).image; - if (src) { - parts.push({ - type: "image_url", - image_url: { - url: src.startsWith("data:") - ? src - : `data:image/png;base64,${src}`, - }, - }); - } - } + pushImagePart(part); } } } @@ -539,37 +663,344 @@ function collectImageParts( return parts; } -function toOpenAIMessage(message: RunMessage): { - role: "system" | "user" | "assistant"; - content: OpenAIMessageContent; -} | null { +function normalizeOpenAIReasoningItem( + value: unknown, +): OpenAIReasoningContentPart | null { + if (!value || typeof value !== "object") { + return null; + } + const item = value as Record; + if (item.type !== "reasoning" || typeof item.id !== "string" || !item.id) { + return null; + } + const summary = Array.isArray(item.summary) + ? item.summary.flatMap((part) => { + if (!part || typeof part !== "object") { + return []; + } + const summaryPart = part as Record; + return summaryPart.type === "summary_text" && + typeof summaryPart.text === "string" + ? [{ type: "summary_text" as const, text: summaryPart.text }] + : []; + }) + : []; + const normalized: OpenAIReasoningContentPart = { + type: "reasoning", + id: item.id, + summary, + }; + if ( + item.status === "in_progress" || + item.status === "completed" || + item.status === "incomplete" + ) { + normalized.status = item.status; + } + return normalized; +} + +function toOpenAIImageEditReferenceMessage( + reference: PendingImageEditReference, +): OpenAIChatMessage | null { + if (!reference.openaiImageGenerationCallId) { + return null; + } + const content: Exclude = []; + const reasoningItem = normalizeOpenAIReasoningItem( + reference.openaiReasoningItem, + ); + if (reasoningItem) { + content.push(reasoningItem); + } + content.push({ + type: "image_generation_call", + id: reference.openaiImageGenerationCallId, + ...(reference.openaiResponseId + ? { response_id: reference.openaiResponseId } + : {}), + }); + return { role: "assistant", content }; +} + +// Refusal flag stamped on assistant metadata when the backend emits the +// `anthropic_refusal` _toolEvent. We drop the refused pair from the next +// request body (Anthropic guidance: leaving refusals in context keeps +// refusing). Metadata (not text) prevents content from spoofing a reset. +function isAnthropicRefusalMessage(message: RunMessage): boolean { + if (message.role !== "assistant") return false; + const metadata = (message as { metadata?: unknown }).metadata as + | { custom?: Record } + | undefined; + return metadata?.custom?.anthropicRefusal === true; +} + +function collectAssistantToolCalls( + message: RunMessage, +): Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; +}> { + const out: Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + }> = []; + for (const part of message.content ?? []) { + if (part.type !== "tool-call") continue; + const tc = part as ToolCallMessagePart & { + argsText?: string; + extra_content?: unknown; + }; + const toolNameLower = (tc.toolName ?? "").toLowerCase(); + const argsObj = + tc.args && typeof tc.args === "object" + ? (tc.args as Record) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + const hasServerToolMarker = Boolean( + argsObj && (argsObj as Record)._server_tool === true, + ); + const isServerSideBuiltin = isServerSideBuiltinToolPart( + toolNameLower, + argsObj, + hasServerToolMarker, + hasNativePart, + ); + if (isServerSideBuiltin) { + // Gemini code_execution / image_generation still need to round- + // trip the native_part payload for native replay; drop the rest. + if (!hasNativePart) continue; + } + const argumentsStr = + typeof tc.argsText === "string" && tc.argsText.length > 0 + ? tc.argsText + : JSON.stringify(tc.args ?? {}); + const entry: { + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + } = { + id: tc.toolCallId, + type: "function" as const, + function: { + name: tc.toolName ?? "", + arguments: argumentsStr, + }, + }; + // Promote args.google to extra_content.google so the backend + // native_part replay branch can find it. The backend only inspects + // extra_content, not function.arguments. + if (tc.extra_content !== undefined) { + entry.extra_content = tc.extra_content; + } else if (argsGoogle) { + entry.extra_content = { google: argsGoogle }; + } + out.push(entry); + } + return out; +} + +function collectToolResultMessages( + message: RunMessage, +): Array<{ + role: "tool"; + content: string; + tool_call_id: string; + name?: string; +}> { + const out: Array<{ + role: "tool"; + content: string; + tool_call_id: string; + name?: string; + }> = []; + for (const part of message.content ?? []) { + if (part.type !== "tool-call") continue; + const tc = part as ToolCallMessagePart; + const result = (tc as { result?: unknown }).result; + // Skip provider-side builtins; see isServerSideBuiltinToolPart. + const argsObj = + tc.args && typeof tc.args === "object" + ? (tc.args as Record) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const toolNameLower = (tc.toolName ?? "").toLowerCase(); + const hasServerToolMarker = Boolean( + argsObj && argsObj._server_tool === true, + ); + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + if ( + isServerSideBuiltinToolPart( + toolNameLower, + argsObj, + hasServerToolMarker, + hasNativePart, + ) + ) { + continue; + } + if (result === undefined || result === null) continue; + let content: string; + if (typeof result === "string") { + // Backend ChatMessage validator rejects role="tool" with empty + // content; serialise a sentinel JSON so legitimately empty tool + // outputs still round-trip the follow-up turn to the provider. + content = result.length > 0 ? result : JSON.stringify({ result: "" }); + } else { + try { + content = JSON.stringify(result); + } catch { + content = String(result); + } + } + out.push({ + role: "tool", + content, + tool_call_id: tc.toolCallId, + ...(tc.toolName ? { name: tc.toolName } : {}), + }); + } + return out; +} + +type SerializedMessage = { + role: "system" | "user" | "assistant" | "tool"; + content: OpenAIMessageContent | null; + tool_calls?: Array<{ + id: string; + type: "function"; + function: { name: string; arguments: string }; + extra_content?: unknown; + }>; + tool_call_id?: string; + name?: string; + /** + * Gemini text-part thoughtSignature stashed during streaming on the + * last text MessagePart. Backend reads + * `extra_content.google.thought_signature` and attaches it to the + * matching Gemini text part on the outbound turn. + */ + extra_content?: unknown; +}; + +function collectAssistantTextThoughtSignature( + message: RunMessage, +): string | undefined { + if (!Array.isArray(message.content)) return undefined; + for (let i = message.content.length - 1; i >= 0; i -= 1) { + const part = message.content[i] as { type?: string } & Record< + string, + unknown + >; + if (part?.type !== "text") continue; + const sig = part._google_thought_signature; + if (typeof sig === "string" && sig) return sig; + } + return undefined; +} + +function toOpenAIMessages(message: RunMessage): SerializedMessage[] { if ( message.role !== "system" && message.role !== "user" && message.role !== "assistant" ) { - return null; + return []; } let textContent = collectTextParts(message).join("\n"); - // Strip inline audio base64 from prior assistant messages to avoid - // inflating token counts (e.g. audio-player responses with embedded WAV). if (message.role === "assistant") { textContent = textContent.replace( /data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g, "[audio]", ); + if (isAnthropicRefusalMessage(message)) { + // Prune refused assistant turn from outbound history; the + // rendered transcript still shows the user-visible notice. + return []; + } } const imageParts = collectImageParts(message); - if (imageParts.length > 0) { - return { - role: message.role, - content: [{ type: "text", text: textContent }, ...imageParts], - }; + const toolCalls = + message.role === "assistant" ? collectAssistantToolCalls(message) : []; + const toolResults = + message.role === "assistant" ? collectToolResultMessages(message) : []; + + const base: SerializedMessage = { + role: message.role, + content: + imageParts.length > 0 + ? [{ type: "text", text: textContent }, ...imageParts] + : textContent, + }; + if (toolCalls.length > 0) { + base.tool_calls = toolCalls; + // OpenAI requires content === null on assistant turns whose + // payload is entirely tool_calls (matches the wire shape Gemini + // expects for the next functionCall replay). + if (!textContent && imageParts.length === 0) { + base.content = null; + } + } + if (message.role === "assistant") { + const sig = collectAssistantTextThoughtSignature(message); + if (sig) { + base.extra_content = { google: { thought_signature: sig } }; + } } - return { role: message.role, content: textContent }; + return toolResults.length > 0 ? [base, ...toolResults] : [base]; +} + +// Thin singular wrapper: returns only the first serialized message +// (without tool_calls or tool follow-ups) so the OpenAI image-edit +// replay path can map a thread to flat OpenAI chat messages without +// pulling in tool history. +function toOpenAIMessage(message: RunMessage): { + role: "system" | "user" | "assistant"; + content: OpenAIMessageContent; +} | null { + const serialized = toOpenAIMessages(message); + if (serialized.length === 0) return null; + const first = serialized[0]; + if ( + first.role !== "system" && + first.role !== "user" && + first.role !== "assistant" + ) { + return null; + } + if (first.content === null || first.content === undefined) { + return null; + } + if (typeof first.content === "string" && !first.content) { + return null; + } + return { + role: first.role, + content: first.content as OpenAIMessageContent, + }; } function extractImageBase64(input: string): string | undefined { @@ -1003,17 +1434,52 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // the user switches chats while waiting for model load / auto-load. const resolvedThreadId = (unstable_threadId ?? runtime.activeThreadId) || undefined; + const resolvedThreadKey = resolvedThreadId ?? null; + const pendingImageEditReferenceForRun = runtime.pendingImageEditReference; + const selectedImageEditReference = + (pendingImageEditReferenceForRun?.threadId ?? null) === + resolvedThreadKey + ? pendingImageEditReferenceForRun + : null; + const clearSelectedImageEditReference = () => { + if (!selectedImageEditReference) { + return; + } + const store = useChatRuntimeStore.getState(); + const pending = store.pendingImageEditReference; + if ( + pending?.openaiImageGenerationCallId === + selectedImageEditReference.openaiImageGenerationCallId && + pending.openaiResponseId === + selectedImageEditReference.openaiResponseId && + (pending.threadId ?? null) === + (selectedImageEditReference.threadId ?? null) + ) { + store.clearPendingImageEditReference(); + } + }; // Wait for in-progress model load to finish before inferring if (runtime.modelLoading) { toast.info("Waiting for model to finish loading…"); - await waitForModelReady(abortSignal); + try { + await waitForModelReady(abortSignal); + } catch (error) { + clearSelectedImageEditReference(); + throw error; + } } if (!useChatRuntimeStore.getState().params.checkpoint) { // Auto-load the smallest downloaded model - const { loaded, blockedByTrustRemoteCode } = - await autoLoadSmallestModel(); + let loaded: boolean; + let blockedByTrustRemoteCode: boolean; + try { + ({ loaded, blockedByTrustRemoteCode } = await autoLoadSmallestModel()); + } catch (error) { + clearSelectedImageEditReference(); + throw error; + } if (!loaded) { toast.error( blockedByTrustRemoteCode @@ -1025,6 +1491,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { : "Pick a model in the top bar, then retry.", }, ); + clearSelectedImageEditReference(); throw new Error("Load a model first."); } } @@ -1037,6 +1504,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toolsEnabled, codeToolsEnabled, imageToolsEnabled, + mcpEnabledForChat, + webFetchToolsEnabled, } = runtime; const externalSelection = parseExternalModelId(params.checkpoint); const isExternalRequest = externalSelection !== null; @@ -1048,6 +1517,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { description: "Turn on Enable connections in Settings → Connections to use hosted models.", }); + clearSelectedImageEditReference(); throw new Error("Connections disabled."); } const externalProvider = isExternalRequest @@ -1063,51 +1533,33 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toast.error("Connection not found.", { description: "Open Settings → Connections and add it again.", }); + clearSelectedImageEditReference(); throw new Error("Connection not found."); } - // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers. + // Local providers and custom Gemini bases allow an empty key. const externalProviderIsCustom = externalProvider ? isCustomProviderType(externalProvider.providerType) : false; - if (isExternalRequest && !externalApiKey && !externalProviderIsCustom) { + const externalProviderIsGeminiCustomBase = Boolean( + externalProvider && + externalProvider.providerType === "gemini" && + isGeminiCustomOpenAICompatBase(externalProvider.baseUrl), + ); + if ( + isExternalRequest && + !externalApiKey && + !externalProviderIsCustom && + !externalProviderIsGeminiCustomBase + ) { toast.error("Missing API key for selected connection.", { description: "Open Settings → Connections and set the API key again.", }); + clearSelectedImageEditReference(); throw new Error("Missing connection API key."); } - const webSearchEnabledForThisTurn = Boolean( - externalProvider && - toolsEnabled && - providerSupportsBuiltinWebSearch(externalProvider.providerType), - ); - const codeExecEnabledForThisTurn = Boolean( - externalProvider && - externalSelection && - codeToolsEnabled && - providerSupportsBuiltinCodeExecution( - externalProvider.providerType, - externalSelection.modelId, - externalProvider.baseUrl, - ), - ); - // web_fetch shares the Search pill with web_search (no separate - // UI toggle), so it follows toolsEnabled. Anthropic is the only - // provider that ships it today; on others providerSupportsBuiltinWebFetch - // returns false and this stays inert. - const webFetchEnabledForThisTurn = Boolean( - externalProvider && - toolsEnabled && - providerSupportsBuiltinWebFetch(externalProvider.providerType), - ); - const providerShipsWebFetch = Boolean( - externalProvider && - providerSupportsBuiltinWebFetch(externalProvider.providerType), - ); - // OpenAI Responses-API image_generation server tool. Pill is - // gated on OpenAI cloud + a Responses-API model id; the backend - // additionally re-checks is_openai_cloud before appending - // {type:"image_generation"} to the request tools array. + // Image-generation flag (OpenAI cloud + Responses-capable model). + // Computed first so Gemini image mode can suppress Search/Code. const imageGenerationEnabledForThisTurn = Boolean( externalProvider && externalSelection && @@ -1118,12 +1570,106 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider.baseUrl, ), ); + // Per-model Search/Code allowances live in + // providerSupportsBuiltin*; this flag just signals image-mode. + const geminiImageModeForThisTurn = + externalProvider?.providerType === "gemini" && + imageGenerationEnabledForThisTurn; + const webSearchEnabledForThisTurn = Boolean( + externalProvider && + externalSelection && + toolsEnabled && + providerSupportsBuiltinWebSearch( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), + ); + const codeExecEnabledForThisTurn = Boolean( + externalProvider && + externalSelection && + codeToolsEnabled && + !geminiImageModeForThisTurn && + providerSupportsBuiltinCodeExecution( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), + ); + // Fetch pill is independent of Search (Anthropic bills web_fetch + // separately from web_search). Sourced from `webFetchToolsEnabled`; + // on providers without web_fetch the toggle is forced off in + // chat-page's runtime setState. + const webFetchEnabledForThisTurn = Boolean( + externalProvider && + webFetchToolsEnabled && + providerSupportsBuiltinWebFetch(externalProvider.providerType), + ); + const providerShipsWebFetch = Boolean( + externalProvider && + providerSupportsBuiltinWebFetch(externalProvider.providerType), + ); - const outboundMessages = messages - .map(toOpenAIMessage) + if (selectedImageEditReference && !imageGenerationEnabledForThisTurn) { + clearSelectedImageEditReference(); + toast.error("Image editing is unavailable", { + description: + "Select an OpenAI image-generation model, then retry the edit.", + }); + throw new Error("Image generation edit unavailable."); + } + + // Drop refused assistant turns + their triggering user prompt; + // otherwise context re-triggers the classifier. + const survivingMessages: RunMessage[] = []; + for (const message of messages) { + if (isAnthropicRefusalMessage(message)) { + const last = survivingMessages.at(-1); + if (last && last.role === "user") { + survivingMessages.pop(); + } + continue; + } + survivingMessages.push(message); + } + + // toOpenAIMessages emits assistant tool_calls + role="tool" + // follow-ups; the backend Gemini translator rebuilds the + // functionCall/functionResponse parts (with thoughtSignature). + const outboundMessages = survivingMessages + .flatMap(toOpenAIMessages) .filter((message): message is NonNullable => Boolean(message), ); + if (selectedImageEditReference) { + const referenceMessage = toOpenAIImageEditReferenceMessage( + selectedImageEditReference, + ); + if (!referenceMessage) { + clearSelectedImageEditReference(); + toast.error("This generated image cannot be edited", { + description: + "The original image reference is missing. Generate the image again, then retry the edit.", + }); + throw new Error("Generated image edit reference missing."); + } + let insertAt = outboundMessages.length; + for (let i = outboundMessages.length - 1; i >= 0; i -= 1) { + if (outboundMessages[i]?.role === "user") { + insertAt = i; + break; + } + } + // OpenAIChatMessage is a structural superset of SerializedMessage + // for the role/content axis the outbound pipeline consumes; cast + // through unknown since referenceMessage carries no tool_calls + // (the image_edit reference is a plain assistant turn). + outboundMessages.splice( + insertAt, + 0, + referenceMessage as unknown as SerializedMessage, + ); + } // Temporary debug toggle: when false, the pre-fetch path is skipped // entirely so retrieval only happens via the LLM-invoked @@ -1244,24 +1790,51 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const webLabel = providerShipsWebFetch ? "web search or web fetch" : "web search"; - if (!webSearchEnabledForThisTurn && !codeExecEnabledForThisTurn) { + // Treat search and fetch as a single "any web tool" axis so + // the guard only warns when neither pill is on; checking + // webSearchEnabledForThisTurn alone mis-fired when only Fetch + // was on and suppressed live web_fetch calls. + const anyWebEnabledForThisTurn = + webSearchEnabledForThisTurn || webFetchEnabledForThisTurn; + if ( + !anyWebEnabledForThisTurn && + !codeExecEnabledForThisTurn && + !imageGenerationEnabledForThisTurn + ) { disabledToolGuard = - `You do not have ${webLabel} or code execution tools in this conversation. ` + + `You do not have ${webLabel}, code execution, or image generation tools in this conversation. ` + "Answer from your own knowledge. " + - "If a request genuinely requires tool use, live data fetch or running code, " + + "If a request genuinely requires tool use, live data fetch, running code, or image generation, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; - } else if (!webSearchEnabledForThisTurn) { + } else if (!anyWebEnabledForThisTurn && !codeExecEnabledForThisTurn) { + disabledToolGuard = + `You do not have ${webLabel} or code execution tools in this conversation. ` + + "You may still use image generation tools when they are available and useful. " + + "If a request genuinely requires live data fetch or running code, " + + "inform the user that you do not have access to these capabilities. " + + "Do not return tool-call syntax inside your response."; + } else if (!anyWebEnabledForThisTurn) { + const availableTools = [ + codeExecEnabledForThisTurn ? "code execution" : null, + imageGenerationEnabledForThisTurn ? "image generation" : null, + ].filter(Boolean); disabledToolGuard = `You do not have ${webLabel} tools in this conversation. ` + - "You may still use code execution tools when they are available and useful. " + + (availableTools.length > 0 + ? `You may still use ${availableTools.join(" and ")} tools when they are available and useful. ` + : "") + "If a request genuinely requires live data fetch or web search tool use, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; } else if (!codeExecEnabledForThisTurn) { + const availableTools = [ + webLabel, + imageGenerationEnabledForThisTurn ? "image generation" : null, + ].filter(Boolean); disabledToolGuard = "You do not have code execution tools in this conversation. " + - `You may still use ${webLabel} tools when they are available and useful. ` + + `You may still use ${availableTools.join(" and ")} tools when they are available and useful. ` + "If a request genuinely requires running code or code execution tool use, " + "inform the user that you do not have access to these capabilities. " + "Do not return tool-call syntax inside your response."; @@ -1279,7 +1852,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { outboundMessages[0] = { ...firstMessage, content: [ - ...firstMessage.content, + ...(Array.isArray(firstMessage.content) + ? firstMessage.content + : []), { type: "text", text: `\n\n${disabledToolGuard}` }, ], }; @@ -1291,8 +1866,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }); } } - const imageBase64 = findLatestUserImageBase64(messages); - const audioBase64 = findLatestUserAudioBase64(messages); + // Scan post-prune history so a refused user turn's image/audio + // doesn't gate or mis-attribute the next non-refused turn. + const imageBase64 = findLatestUserImageBase64(survivingMessages); + const audioBase64 = findLatestUserAudioBase64(survivingMessages); // Block when ANY image is in the outbound payload (current or // prior turns) and the loaded model can't process images. Keeps @@ -1321,6 +1898,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const gatedThreadKey = resolvedThreadId || "__default"; runtime.setThreadRunning(gatedThreadKey, true); runtime.setThreadRunning(gatedThreadKey, false); + clearSelectedImageEditReference(); throw new Error(imageGateReason); } } @@ -1328,7 +1906,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (audioBase64) { const audioName = runtime.pendingAudioName; if (audioName) { - const lastUserMsg = [...messages] + const lastUserMsg = [...survivingMessages] .reverse() .find((m) => m.role === "user"); if (lastUserMsg) sentAudioNames.set(lastUserMsg.id, audioName); @@ -1425,20 +2003,56 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let cumulativeText = ""; let reasoningStartAt: number | null = null; let reasoningDuration = 0; - // Tracks whether we are currently inside a `` block opened by - // a `delta.reasoning_content` chunk. Kimi (kimi-k2.6, kimi-k2-thinking) - // and DeepSeek's reasoner stream their thinking as a separate - // `reasoning_content` field on the chat-completion delta — not as - // `content`, not as a structured part. We wrap those chunks with - // inline `...` so the existing parseAssistantContent - // lifts them into the reasoning panel the same way it does for - // local Harmony models. State has to live outside the SSE loop - // because the close tag fires when the next chunk carries content - // (or when the stream ends). + // True while wrapping a `delta.reasoning_content` stream in + // ... for parseAssistantContent. Lives outside + // the SSE loop because the close tag fires when content arrives. let reasoningContentOpen = false; - // Tool call content parts — accumulated and yielded cumulatively. - // result is set directly on the tool-call part when tool_end arrives. + // Tool call parts, cumulative; result lands on tool_end. const toolCallParts: ToolCallMessagePart[] = []; + // Latest Gemini text-part thoughtSignature; pinned onto the final + // text MessagePart so next-turn replay carries it. + let latestTextThoughtSignature: string | undefined; + const pinTextThoughtSignature = ( + parts: T[], + ): T[] => { + if (!latestTextThoughtSignature || parts.length === 0) return parts; + for (let i = parts.length - 1; i >= 0; i -= 1) { + if (parts[i].type === "text") { + parts[i] = { + ...parts[i], + _google_thought_signature: latestTextThoughtSignature, + } as T; + break; + } + } + return parts; + }; + const orderAssistantContent = ( + textParts: ReturnType, + ) => { + const imageToolParts = toolCallParts.filter( + (part) => part.toolName === "image_generation", + ); + const otherToolParts = toolCallParts.filter( + (part) => part.toolName !== "image_generation", + ); + return [...otherToolParts, ...textParts, ...imageToolParts]; + }; + // Anthropic document_citations tool_event payload, converted to + // Sources-panel source parts at end-of-stream so the inline [N] + // markers have matching entries. + const documentCitationParts: Array<{ + type: "source"; + sourceType: "url"; + id: string; + url: string; + title: string; + metadata?: { description: string }; + }> = []; + // Latched on the `anthropic_refusal` tool event; stamped onto the + // final assistant metadata as `custom.anthropicRefusal` to drive + // the history-prune above. + let anthropicRefusalSeen = false; let serverMetadata: { usage?: ServerUsage; timings?: ServerTimings; @@ -1504,6 +2118,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { { isReasoningProvider: externalProvider.isReasoningModel === true, + baseUrl: externalProvider.baseUrl ?? null, }, ) : { @@ -1536,13 +2151,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { forceRefreshPublicKey = false, ): Promise => { if (externalSelection && externalProvider) { - // OpenAI shell-tool container reuse: pull the per-thread - // container_id (if any) so subsequent turns in the same - // thread reference the existing container instead of - // auto-creating a fresh one. Empty string / undefined → - // backend falls back to container_auto. Anthropic uses - // the parallel `anthropicCodeExecContainerId` field below - // (sent as `container` on /v1/messages). + // Per-thread container reuse; empty/undefined falls back to + // container_auto. Anthropic uses anthropicCodeExecContainerId. let openaiCodeExecContainerId: string | null = null; let anthropicCodeExecContainerId: string | null = null; if (codeExecEnabledForThisTurn && resolvedThreadId) { @@ -1556,17 +2166,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = null; anthropicCodeExecContainerId = null; } - // Pre-send container validation (OpenAI only). The list - // endpoint already filters status==="expired" server-side - // (studio/backend/routes/inference.py — list_openai_containers), - // so membership in this set means "OpenAI will accept it - // as container_reference". A stale id silently dropped here - // falls through to the inheritance + lazy-create logic - // below, so the user never sees "Container is expired" in - // the chat thread. On list-call failure we leave - // activeContainerIds null and skip validation — the - // backend's transparent retry path is the safety net for - // that case. + // Pre-send container validation (OpenAI). Stale ids drop + // silently and fall through to lazy-create. On list-call + // failure, skip and rely on the backend's retry path. let activeContainerIds: Set | null = null; if (externalProvider.providerType === "openai") { try { @@ -1589,15 +2191,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { openaiCodeExecContainerId = null; } } - // Cross-thread inheritance: when the active thread has - // no container yet, default to the one most recently - // used on *any* other thread (provider-scoped). - // Matches what the Code Execution settings section - // shows in the picker, and keeps the user from getting - // a fresh container on every new thread. The picker - // can still be set to "Auto-create per thread" - // explicitly to opt into a fresh container — but - // that's done via the dropdown, not silently. + // Cross-thread inheritance: reuse the most recently used + // container from any other thread; opt-out via the picker. if ( !openaiCodeExecContainerId && externalProvider.providerType === "openai" @@ -1630,13 +2225,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { /* fall through to lazy-create below */ } } - // Lazy pre-create when there's no inherited container. - // We always POST /v1/containers ourselves (rather than - // letting the backend send container_auto) so every - // container shows up in the picker with a friendly - // English-word name and the user's configured TTL. - // Falls back to container_auto only if the POST fails - // — keeps the chat moving in that case. + // Pre-create our own container (vs container_auto) so it + // shows up in the picker with a friendly name and the + // configured TTL. Falls back to container_auto on failure. if ( !openaiCodeExecContainerId && externalProvider.providerType === "openai" @@ -1683,35 +2274,25 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ...(externalCapabilities?.topP !== false ? { top_p: params.topP } : {}), - // Clamp to the cross-provider output cap so a maxTokens value - // carried over from a local-model session does not blow past - // provider limits (e.g. Claude Opus 400s on >128k). Also - // floor to the provider's documented minimum — Kimi's - // thinking models need >=16k or the response truncates - // before the answer fits alongside reasoning_content. + // Floor at the provider's documented min (Kimi thinking + // needs >=16k); clamp at the per-model max. max_tokens: Math.min( Math.max( params.maxTokens, getExternalMinOutputTokens(externalProvider?.providerType), ), - EXTERNAL_MAX_OUTPUT_TOKENS, + getExternalMaxOutputTokens( + externalProvider?.providerType, + externalSelection?.modelId, + ), ), - // Only forward sampling knobs the provider actually accepts; the - // backend's external-provider proxy is param-permissive and would - // surface a 400 from providers that reject unknown fields (e.g. - // OpenAI rejects top_k, Anthropic/DeepSeek reject presence_penalty). + // Forward only sampling knobs the provider accepts. ...(externalCapabilities?.topK ? { top_k: params.topK } : {}), ...(externalCapabilities?.presencePenalty ? { presence_penalty: params.presencePenalty } : {}), - // Built-in tools: Search pill maps to provider-side - // web_search (currently OpenAI / Anthropic / OpenRouter / - // Kimi); Code pill maps to Anthropic's server-side - // code_execution_20250825 tool (Anthropic is the only - // external provider that ships one today). Backend - // translates enabled_tools into each provider's tool - // schema — for Anthropic that's the entries appended to - // body["tools"] inside _stream_anthropic. + // Compose the enabled_tools list from the active pills; + // backend maps each name to the provider's tool schema. ...(webSearchEnabledForThisTurn || webFetchEnabledForThisTurn || codeExecEnabledForThisTurn || @@ -1720,19 +2301,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { enable_tools: true, enabled_tools: [ ...(webSearchEnabledForThisTurn ? ["web_search"] : []), - // Pair web_fetch with the Search pill on any - // provider that ships it (Anthropic today). The - // common workflow is "search returns URLs, fetch - // reads them"; without web_fetch the model can - // surface a citation but cannot quote from the - // page body, which is the whole point of the - // tool. There is no separate UI toggle yet. ...(webFetchEnabledForThisTurn ? ["web_fetch"] : []), ...(codeExecEnabledForThisTurn ? ["code_execution"] : []), - // OpenAI Responses-API only: `image_generation` - // returns inline image_generation_call output - // items; the backend's _stream_openai_responses - // path translates them to assistant tool events. ...(imageGenerationEnabledForThisTurn ? ["image_generation"] : []), @@ -1781,6 +2351,16 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { isPromptCacheTtl(externalProvider.promptCacheTtl) ? { prompt_cache_ttl: externalProvider.promptCacheTtl } : {}), + // Anthropic fast mode (Opus 4.6 / 4.7 only); backend + // silently drops on unsupported models as a second + // line of defence. + ...(params.fastMode && + providerSupportsFastMode( + externalProvider.providerType, + externalSelection.modelId, + ) + ? { fast_mode: true } + : {}), ...(externalReasoningCaps.supportsReasoning ? externalReasoningCaps.reasoningStyle === "reasoning_effort" ? externalReasoningEnabled @@ -1822,7 +2402,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { ? { preserve_thinking: preserveThinking } : {}), ...(supportsTools && - (toolsEnabled || codeToolsEnabled || ragToolPathTaken) + (toolsEnabled || + codeToolsEnabled || + ragToolPathTaken || + mcpEnabledForChat) ? { enable_tools: true, enabled_tools: [ @@ -1849,6 +2432,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }, } : {}), + mcp_enabled: mcpEnabledForChat, auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, max_tool_calls_per_message: @@ -1865,10 +2449,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let retriedWithRefreshedKey = false; while (true) { try { - const stream = streamChatCompletions( - await buildRequestPayload(retriedWithRefreshedKey), - abortSignal, - ); + let requestPayload: OpenAIChatCompletionsRequest; + try { + requestPayload = await buildRequestPayload(retriedWithRefreshedKey); + } catch (error) { + clearSelectedImageEditReference(); + throw error; + } + clearSelectedImageEditReference(); + const stream = streamChatCompletions(requestPayload, abortSignal); for await (const chunk of stream) { // Handle tool status events @@ -1887,11 +2476,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { chunk as unknown as { _toolEvent?: Record } )._toolEvent; if (toolEvent !== undefined) { - // OpenAI shell-tool container persistence — see - // ThreadRecord.openaiCodeExecContainerId. The backend - // emits these synthetic events on the OpenAI Responses - // SSE stream after capturing the container_id from a - // response, or detecting an expired-container error. + // Persist container_id onto the thread (OpenAI / Anthropic). if (toolEvent.type === "container_ready") { const newContainerId = toolEvent.container_id as | string @@ -1907,6 +2492,27 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } continue; } + if (toolEvent.type === "document_citations") { + // Convert Anthropic citations_delta footnotes into + // Sources-panel entries matching the inline [N] markers. + const cits = toolEvent.citations; + if (Array.isArray(cits)) { + cits.forEach((entry, idx) => { + if (!entry || typeof entry !== "object") return; + const part = documentCitationToSource( + entry as Record, + idx, + ); + if ( + part && + !documentCitationParts.some((p) => p.id === part.id) + ) { + documentCitationParts.push(part); + } + }); + } + continue; + } if (toolEvent.type === "container_invalidated") { if (resolvedThreadId) { const field = @@ -1919,6 +2525,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } continue; } + if (toolEvent.type === "anthropic_refusal") { + // Latch the backend refusal signal so the final + // message metadata can drive the prune. + anthropicRefusalSeen = true; + continue; + } if (toolEvent.type === "tool_start") { const id = (toolEvent.tool_call_id as string) || @@ -1953,6 +2565,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { size?: string; quality?: string; background?: string; + prompt?: string; }; const imageB64 = toolEvent.image_b64 as string | undefined; if ( @@ -1960,12 +2573,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { typeof imageB64 === "string" && imageB64 ) { - // OpenAI Responses image_generation_call: the - // backend stashes the base64 PNG/WebP/JPEG on - // separate `image_b64` / `image_mime` fields on - // the synthetic _toolEvent so the JSON result - // string stays small enough to log. Repackage as - // a structured result for the dedicated tool UI. + // Backend keeps base64 on separate image_b64 / + // image_mime fields so logs stay small; repackage. parsedResult = { image_b64: imageB64, image_mime: @@ -1974,6 +2583,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { size: toolEvent.size as string | undefined, quality: toolEvent.quality as string | undefined, background: toolEvent.background as string | undefined, + prompt: toolEvent.prompt as string | undefined, }; } else if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); @@ -1991,16 +2601,106 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } else { parsedResult = rawResult; } + // Merge tool_end args first, then Gemini native_part. + const nextArgs = + toolEvent.arguments && + typeof toolEvent.arguments === "object" + ? (toolEvent.arguments as ToolCallMessagePart["args"]) + : undefined; + const mergedArgs: ToolCallMessagePart["args"] = { + ...(toolCallParts[idx].args ?? {}), + ...(nextArgs ?? {}), + } as ToolCallMessagePart["args"]; + // Merge tool_end native_part into args.google so the + // outbound translator replays both start (executableCode) + // and end (result / inlineData) on the same turn. + // Concatenate parts so each keeps its own thoughtSignature. + const endGoogle = ( + toolEvent as { google?: { native_part?: unknown } } + ).google; + if ( + endGoogle && + typeof endGoogle === "object" && + endGoogle.native_part && + typeof endGoogle.native_part === "object" + ) { + const argsObj = mergedArgs as Record; + const existingGoogle = (argsObj.google ?? {}) as Record< + string, + unknown + >; + const existingNative = + (existingGoogle.native_part as Record< + string, + unknown + >) ?? {}; + const endNative = endGoogle.native_part as Record< + string, + unknown + >; + // Extract part entries from either parts:[...] or + // legacy single-object native_part. Legacy + // thoughtSignature always belongs on executableCode. + const collectParts = ( + native: Record, + ): Record[] => { + if (Array.isArray(native.parts)) { + return (native.parts as unknown[]).filter( + (entry): entry is Record => + Boolean(entry) && + typeof entry === "object" && + !Array.isArray(entry), + ); + } + const out: Record[] = []; + const legacySig = + typeof native.thoughtSignature === "string" + ? native.thoughtSignature + : typeof native.thought_signature === "string" + ? (native.thought_signature as string) + : null; + for (const key of [ + "executableCode", + "codeExecutionResult", + "inlineData", + ] as const) { + const sub = native[key]; + if (sub && typeof sub === "object") { + const entry: Record = { + [key]: sub, + }; + if (key === "executableCode" && legacySig) { + entry.thoughtSignature = legacySig; + } + out.push(entry); + } + } + return out; + }; + const mergedParts = [ + ...collectParts(existingNative), + ...collectParts(endNative), + ]; + argsObj.google = { + ...existingGoogle, + native_part: { parts: mergedParts }, + }; + } toolCallParts[idx] = { ...toolCallParts[idx], + args: mergedArgs, + argsText: JSON.stringify(mergedArgs ?? {}), result: parsedResult, }; } } - // Yield cumulative state so tool UI updates (tools first, text after) - const textParts = parseAssistantContent(cumulativeText); + // Cumulative yield. orderAssistantContent puts search/ + // code before text and generated images after. + const textParts = pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ); yield { - content: [...toolCallParts, ...textParts], + content: orderAssistantContent(textParts), metadata: { timing: buildTiming( streamStartTime, @@ -2025,11 +2725,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } totalChunks += 1; - // OpenRouter's free router (openrouter/free) picks a different - // underlying free model per request and reports it in every - // chunk's top-level `model` field. Latch the first non-empty - // value that differs from the requested checkpoint so the - // header chip can render "openrouter/free:". + // Latch the chunk's `model` field so the openrouter/free + // chip can show the chosen underlying model. if ( isExternalRequest && externalProvider?.providerType === "openrouter" && @@ -2048,30 +2745,37 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } } const rawDelta = chunk.choices?.[0]?.delta?.content; - // Providers like Mistral's magistral return delta.content as an - // array of structured parts; normalize to text (with thinking - // parts re-wrapped as inline tags) so the rest of the - // accumulator stays string-based. + // Normalize structured delta.content (mistral magistral) to text. const delta = extractDeltaText(rawDelta); - // Kimi (kimi-k2.6, kimi-k2-thinking) and DeepSeek reasoner - // stream thinking via `delta.reasoning_content` as a plain - // string field — separate from `delta.content` which carries - // the answer. Wrap reasoning chunks inline as ... - // so parseAssistantContent treats them like any - // other reasoning. The close tag fires when the next chunk - // brings content, or when the stream ends. + // Latest Gemini text-part thoughtSignature for next-turn replay. + const deltaExtraContent = ( + chunk.choices?.[0]?.delta as + | { extra_content?: unknown } + | undefined + )?.extra_content; + if ( + deltaExtraContent && + typeof deltaExtraContent === "object" + ) { + const eGoogle = (deltaExtraContent as Record) + .google; + if (eGoogle && typeof eGoogle === "object") { + const sig = (eGoogle as Record) + .thought_signature; + if (typeof sig === "string" && sig) { + latestTextThoughtSignature = sig; + } + } + } + // Kimi / DeepSeek stream thinking via delta.reasoning_content. + // Wrap inline as ... for parseAssistantContent. const rawReasoning = ( chunk.choices?.[0]?.delta as | { reasoning_content?: unknown } | undefined )?.reasoning_content; - // OpenRouter uses a third reasoning shape: a structured - // `delta.reasoning_details` array of parts (each carrying - // `text`). The router emits this regardless of which - // underlying provider it picked, so we extract here and - // merge into the same ... wrap path used - // for Kimi / DeepSeek reasoning_content. See - // https://openrouter.ai/docs/guides/best-practices/reasoning-tokens + // OpenRouter ships reasoning as delta.reasoning_details[] + // regardless of underlying provider; merge into the same wrap path. const rawReasoningDetails = ( chunk.choices?.[0]?.delta as | { reasoning_details?: unknown } @@ -2089,6 +2793,138 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const reasoning = (typeof rawReasoning === "string" ? rawReasoning : "") + reasoningFromDetails; + // OpenAI delta.tool_calls: streams fragments by index; + // accumulate into one part. extra_content carries Gemini 3 + // thoughtSignature for next-turn replay. + const rawDeltaToolCalls = ( + chunk.choices?.[0]?.delta as + | { tool_calls?: unknown } + | undefined + )?.tool_calls; + if ( + Array.isArray(rawDeltaToolCalls) && + rawDeltaToolCalls.length > 0 + ) { + for (const tc of rawDeltaToolCalls) { + if (!tc || typeof tc !== "object") continue; + const call = tc as { + id?: string; + index?: number; + function?: { name?: string; arguments?: string }; + extra_content?: unknown; + }; + const idx = + typeof call.index === "number" ? call.index : undefined; + const stableId = call.id; + // Match an existing fragment by id first (canonical), + // then by index slot. Fall back to a freshly-minted + // tool_call_ id for streams that send neither. + let existing = stableId + ? toolCallParts.find((p) => p.toolCallId === stableId) + : undefined; + if (!existing && idx !== undefined) { + existing = toolCallParts.find( + (p) => + ( + p as ToolCallMessagePart & { _delta_index?: number } + )._delta_index === idx, + ); + } + const argsFragment = call.function?.arguments ?? ""; + if (existing) { + const prevName = existing.toolName ?? ""; + const nextName = call.function?.name ?? prevName; + const merged = + (existing.argsText ?? "") + argsFragment; + let parsedArgs: + ToolCallMessagePart["args"] = existing.args ?? {}; + if (merged) { + try { + parsedArgs = JSON.parse( + merged, + ) as ToolCallMessagePart["args"]; + } catch { + parsedArgs = { + _raw: merged, + } as ToolCallMessagePart["args"]; + } + } + const prevExtra = ( + existing as ToolCallMessagePart & { + extra_content?: unknown; + } + ).extra_content; + const updated: ToolCallMessagePart & { + _delta_index?: number; + extra_content?: unknown; + } = { + ...(existing as ToolCallMessagePart), + toolName: nextName, + argsText: merged, + args: parsedArgs, + ...(call.extra_content !== undefined + ? { extra_content: call.extra_content } + : prevExtra !== undefined + ? { extra_content: prevExtra } + : {}), + ...(idx !== undefined ? { _delta_index: idx } : {}), + }; + const replaceIdx = toolCallParts.indexOf(existing); + if (replaceIdx >= 0) { + toolCallParts[replaceIdx] = updated; + } + } else { + const callId = + stableId || + `tool_call_${idx ?? toolCallParts.length}`; + const argsText = argsFragment; + let parsedArgs: ToolCallMessagePart["args"] = {}; + if (argsText) { + try { + parsedArgs = JSON.parse( + argsText, + ) as ToolCallMessagePart["args"]; + } catch { + parsedArgs = { + _raw: argsText, + } as ToolCallMessagePart["args"]; + } + } + const fresh: ToolCallMessagePart & { + _delta_index?: number; + extra_content?: unknown; + } = { + type: "tool-call" as const, + toolCallId: callId, + toolName: call.function?.name ?? "", + argsText, + args: parsedArgs, + ...(call.extra_content !== undefined + ? { extra_content: call.extra_content } + : {}), + ...(idx !== undefined ? { _delta_index: idx } : {}), + }; + toolCallParts.push(fresh); + } + } + yield { + content: [ + ...toolCallParts, + ...pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ), + ], + metadata: { + timing: buildTiming( + streamStartTime, + totalChunks, + firstTokenTime, + ), + custom: { reasoningDuration }, + }, + }; + continue; + } if (!delta && !reasoning) { continue; } @@ -2114,21 +2950,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } cumulativeText += delta; } - // Mistral's magistral occasionally emits a trailing - // template-literal artifact (e.g. "${response}") at the end of - // an otherwise complete answer. It is never part of a real - // reply, so strip a trailing `${...}` token from external - // provider streams. The regex anchors to end-of-string and is - // idempotent — fragments mid-stream (e.g. "${re") leave the - // string untouched and only collapse once the closing brace - // arrives. Local-model output is left alone. + // Strip a trailing ${...} template-literal artifact from + // external streams (mistral magistral occasionally emits one). if (isExternalRequest) { cumulativeText = cumulativeText.replace( /\s*\$\{[^}]*\}\s*$/, "", ); } - const parts = parseAssistantContent(cumulativeText); + const parts = pinTextThoughtSignature( + parseAssistantContent(cumulativeText), + ); if ( parts.some((part) => part.type === "reasoning") && @@ -2148,7 +2980,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (parts.length > 0 || toolCallParts.length > 0) { yield { - content: [...toolCallParts, ...parts], + content: orderAssistantContent(parts), metadata: { timing: buildTiming( streamStartTime, @@ -2231,18 +3063,31 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const finalTokPerSec = meta?.timings?.predicted_per_second; const serverPromptEvalTime = meta?.timings?.prompt_ms; - // Update context usage in store if we got valid server data + // Prefer llama-server timings; fall back to provider usage envelope. + const cachedTokens = + meta?.timings?.cache_n ?? + meta?.usage?.prompt_tokens_details?.cached_tokens ?? + meta?.usage?.cache_read_input_tokens ?? + 0; + // Anthropic-only (billed at the write premium). + const cacheWriteTokens = meta?.usage?.cache_creation_input_tokens ?? 0; + + // Gate on the captured checkpoint still being active so a late + // completion from provider A doesn't populate the bar after the + // user switched to provider B mid-stream. if ( meta?.usage && typeof meta.usage.prompt_tokens === "number" && typeof meta.usage.completion_tokens === "number" && - typeof meta.usage.total_tokens === "number" + typeof meta.usage.total_tokens === "number" && + useChatRuntimeStore.getState().params.checkpoint === params.checkpoint ) { useChatRuntimeStore.getState().setContextUsage({ promptTokens: meta.usage.prompt_tokens, completionTokens: meta.usage.completion_tokens, totalTokens: meta.usage.total_tokens, - cachedTokens: meta.timings?.cache_n ?? 0, + cachedTokens, + cacheWriteTokens, }); } @@ -2258,21 +3103,26 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { yield { content: [ - ...toolCallParts, - ...parseAssistantContent(cumulativeText), + ...orderAssistantContent( + pinTextThoughtSignature(parseAssistantContent(cumulativeText)), + ), ...sourceParts, + ...documentCitationParts, ], metadata: { timing: finalTiming, custom: { reasoningDuration, + // Persisted refusal flag driving the two-pass prune. + anthropicRefusal: anthropicRefusalSeen || undefined, serverTimings: meta?.timings ?? undefined, contextUsage: meta?.usage ? { promptTokens: meta.usage.prompt_tokens, completionTokens: meta.usage.completion_tokens, totalTokens: meta.usage.total_tokens, - cachedTokens: meta.timings?.cache_n ?? 0, + cachedTokens, + cacheWriteTokens, modelId: params.checkpoint, } : undefined, diff --git a/studio/frontend/src/features/chat/api/mcp-servers-api.ts b/studio/frontend/src/features/chat/api/mcp-servers-api.ts new file mode 100644 index 0000000000..be88d12664 --- /dev/null +++ b/studio/frontend/src/features/chat/api/mcp-servers-api.ts @@ -0,0 +1,116 @@ +// 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 { authFetch } from "@/features/auth"; +import { formatFastApiDetail } from "@/lib/format-fastapi-error"; + +export interface McpServerConfig { + id: string; + display_name: string; + url: string; + headers: Record; + is_enabled: boolean; + use_oauth: boolean; + created_at: string; + updated_at: string; +} + +export interface McpServerProbeResult { + ok: boolean; + tool_count: number; + error: string | null; +} + +function parseErrorText(status: number, body: unknown): string { + if (body && typeof body === "object") { + const { detail, message } = body as { detail?: unknown; message?: unknown }; + const formatted = formatFastApiDetail(detail); + if (formatted) return formatted; + if (typeof message === "string" && message) return message; + } + return `Request failed (${status})`; +} + +async function mcpRequest( + path: string, + init?: { method?: string; body?: object }, +): Promise { + const response = await authFetch(`/api/mcp/servers${path}`, { + method: init?.method, + headers: init?.body ? { "Content-Type": "application/json" } : undefined, + body: init?.body ? JSON.stringify(init.body) : undefined, + }); + // 204 No Content (DELETE) has no body — calling .json() would throw. + if (response.status === 204) return undefined as T; + const json = await response.json().catch(() => null); + if (!response.ok) throw new Error(parseErrorText(response.status, json)); + return json as T; +} + +export function listMcpServers(): Promise { + return mcpRequest("/"); +} + +export function createMcpServer(payload: { + displayName: string; + url: string; + headers?: Record; + isEnabled?: boolean; + useOauth?: boolean; +}): Promise { + return mcpRequest("/", { + method: "POST", + body: { + display_name: payload.displayName, + url: payload.url, + headers: payload.headers ?? null, + is_enabled: payload.isEnabled ?? true, + use_oauth: payload.useOauth ?? false, + }, + }); +} + +export function updateMcpServer( + serverId: string, + payload: { + displayName?: string; + url?: string; + /** null = drop stored headers; omit to leave as-is */ + headers?: Record | null; + isEnabled?: boolean; + useOauth?: boolean; + }, +): Promise { + const body: Record = {}; + if (payload.displayName !== undefined) body.display_name = payload.displayName; + if (payload.url !== undefined) body.url = payload.url; + if (payload.headers !== undefined) body.headers = payload.headers; + if (payload.isEnabled !== undefined) body.is_enabled = payload.isEnabled; + if (payload.useOauth !== undefined) body.use_oauth = payload.useOauth; + return mcpRequest(`/${serverId}`, { method: "PUT", body }); +} + +export function deleteMcpServer(serverId: string): Promise { + return mcpRequest(`/${serverId}`, { method: "DELETE" }); +} + +export function refreshMcpServerTools( + serverId: string, +): Promise { + return mcpRequest(`/${serverId}/refresh`, { method: "POST" }); +} + +export function testMcpServer(payload: { + url: string; + headers?: Record; + useOauth?: boolean; +}): Promise { + return mcpRequest("/test", { + method: "POST", + body: { + url: payload.url, + headers: payload.headers ?? null, + use_oauth: payload.useOauth ?? false, + }, + }); +} diff --git a/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx new file mode 100644 index 0000000000..35b5aca64c --- /dev/null +++ b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx @@ -0,0 +1,497 @@ +// 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 { useCallback, useEffect, useState } from "react"; +import { toast } from "sonner"; +import { Delete02Icon, Edit03Icon, PlusSignIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { RefreshCwIcon } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Spinner } from "@/components/ui/spinner"; +import { Switch } from "@/components/ui/switch"; +import { + type McpServerConfig, + createMcpServer, + deleteMcpServer, + listMcpServers, + refreshMcpServerTools, + testMcpServer, + updateMcpServer, +} from "./api/mcp-servers-api"; +type HeaderRow = { id: string; key: string; value: string }; + +type FormState = { + displayName: string; + url: string; + headers: HeaderRow[]; + useOauth: boolean; +}; + +const EMPTY_FORM: FormState = { + displayName: "", + url: "", + headers: [], + useOauth: false, +}; + +function newRowId(): string { + return `r_${Math.random().toString(36).slice(2, 10)}`; +} + +function headersFromObject(headers: Record): HeaderRow[] { + return Object.entries(headers).map(([k, v]) => ({ + id: newRowId(), + key: k, + value: v, + })); +} + +function headersToObject(rows: HeaderRow[]): Record | undefined { + const out: Record = {}; + for (const row of rows) { + const key = row.key.trim(); + if (!key) continue; + out[key] = row.value; + } + return Object.keys(out).length > 0 ? out : undefined; +} + +function isValidUrl(url: string): boolean { + const trimmed = url.trim(); + if (!trimmed) return false; + try { + const parsed = new URL(trimmed); + return parsed.protocol === "http:" || parsed.protocol === "https:"; + } catch { + return false; + } +} + +function HeadersEditor({ + rows, + onChange, +}: { + rows: HeaderRow[]; + onChange: (rows: HeaderRow[]) => void; +}) { + const update = (id: string, patch: Partial) => + onChange(rows.map((row) => (row.id === id ? { ...row, ...patch } : row))); + const add = () => + onChange([...rows, { id: newRowId(), key: "", value: "" }]); + const remove = (id: string) => + onChange(rows.filter((row) => row.id !== id)); + + return ( + <> +
+ + +
+ {rows.length === 0 ? ( +
+ Optional. Add an Authorization header here for servers + that require auth. +
+ ) : ( +
+ {rows.map((row) => ( +
+ update(row.id, { key: e.target.value })} + /> + update(row.id, { value: e.target.value })} + /> + +
+ ))} +
+ )} + + ); +} + +export interface ChatMcpServersDialogProps { + open: boolean; + onOpenChange: (open: boolean) => void; +} + +type View = + | { kind: "list" } + | { kind: "create" } + | { kind: "edit"; id: string }; + +export function ChatMcpServersDialog({ + open, + onOpenChange, +}: ChatMcpServersDialogProps) { + const [servers, setServers] = useState([]); + const [loading, setLoading] = useState(false); + const [view, setView] = useState({ kind: "list" }); + const [form, setForm] = useState(EMPTY_FORM); + const [saving, setSaving] = useState(false); + const [testing, setTesting] = useState(false); + const [refreshingId, setRefreshingId] = useState(null); + + const refresh = useCallback(async () => { + setLoading(true); + try { + const rows = await listMcpServers(); + setServers(rows); + } catch (err) { + toast.error("Failed to load MCP servers", { + description: err instanceof Error ? err.message : String(err), + }); + } finally { + setLoading(false); + } + }, []); + + useEffect(() => { + if (!open) return; + refresh(); + }, [open, refresh]); + + function startCreate() { + setView({ kind: "create" }); + setForm(EMPTY_FORM); + } + + function startEdit(server: McpServerConfig) { + setView({ kind: "edit", id: server.id }); + setForm({ + displayName: server.display_name, + url: server.url, + headers: headersFromObject(server.headers ?? {}), + useOauth: server.use_oauth ?? false, + }); + } + + function cancelForm() { + setView({ kind: "list" }); + setForm(EMPTY_FORM); + } + + async function testConnection() { + const trimmedUrl = form.url.trim(); + if (!isValidUrl(trimmedUrl)) { + toast.error("Enter a valid http:// or https:// URL first"); + return; + } + setTesting(true); + try { + const result = await testMcpServer({ + url: trimmedUrl, + headers: headersToObject(form.headers), + useOauth: form.useOauth, + }); + if (result.ok) { + toast.success( + `Connected (${result.tool_count} tool${result.tool_count === 1 ? "" : "s"})`, + ); + } else { + toast.error("Connection failed", { + description: result.error ?? "Unknown error", + }); + } + } catch (err) { + toast.error("Connection test failed", { + description: err instanceof Error ? err.message : String(err), + }); + } finally { + setTesting(false); + } + } + + async function submitForm() { + const trimmedName = form.displayName.trim(); + const trimmedUrl = form.url.trim(); + if (!trimmedName) { + toast.error("Display name is required"); + return; + } + if (!trimmedUrl) { + toast.error("URL is required"); + return; + } + if (!isValidUrl(trimmedUrl)) { + toast.error("URL must start with http:// or https://"); + return; + } + setSaving(true); + try { + const headers = headersToObject(form.headers); + if (view.kind === "edit") { + await updateMcpServer(view.id, { + displayName: trimmedName, + url: trimmedUrl, + headers: headers ?? null, + useOauth: form.useOauth, + }); + toast.success("MCP server updated"); + } else { + await createMcpServer({ + displayName: trimmedName, + url: trimmedUrl, + headers: headers, + useOauth: form.useOauth, + }); + toast.success("MCP server added"); + } + cancelForm(); + await refresh(); + } catch (err) { + toast.error("Save failed", { + description: err instanceof Error ? err.message : String(err), + }); + } finally { + setSaving(false); + } + } + + async function removeServer(server: McpServerConfig) { + const ok = window.confirm(`Delete MCP server "${server.display_name}"?`); + if (!ok) return; + try { + await deleteMcpServer(server.id); + await refresh(); + } catch (err) { + toast.error("Delete failed", { + description: err instanceof Error ? err.message : String(err), + }); + } + } + + async function toggleEnabled(server: McpServerConfig, next: boolean) { + // Optimistic update so the switch doesn't snap back during the round-trip. + setServers((rows) => + rows.map((row) => + row.id === server.id ? { ...row, is_enabled: next } : row, + ), + ); + try { + await updateMcpServer(server.id, { isEnabled: next }); + } catch (err) { + setServers((rows) => + rows.map((row) => + row.id === server.id ? { ...row, is_enabled: !next } : row, + ), + ); + toast.error("Update failed", { + description: err instanceof Error ? err.message : String(err), + }); + } + } + + async function refreshTools(server: McpServerConfig) { + setRefreshingId(server.id); + try { + const result = await refreshMcpServerTools(server.id); + if (result.ok) { + toast.success( + `Refreshed "${server.display_name}" (${result.tool_count} tool${result.tool_count === 1 ? "" : "s"})`, + ); + } else { + toast.error(`Refresh failed for "${server.display_name}"`, { + description: result.error ?? "Unknown error", + }); + } + } catch (err) { + toast.error("Refresh failed", { + description: err instanceof Error ? err.message : String(err), + }); + } finally { + setRefreshingId(null); + } + } + + const showForm = view.kind !== "list"; + + return ( + + + + MCP Servers + + Register remote MCP servers. + + + + {showForm ? ( +
+
+ + + setForm((prev) => ({ ...prev, displayName: e.target.value })) + } + placeholder="e.g. GitHub MCP" + /> +
+
+ + + setForm((prev) => ({ ...prev, url: e.target.value })) + } + placeholder="https://example.com/mcp" + /> +
+ +
+
+ + + For servers that require browser-based authentication + (GitHub, Linear, etc.). A browser window will open on first + connect. + +
+ + setForm((prev) => ({ ...prev, useOauth })) + } + /> +
+ + setForm((prev) => ({ ...prev, headers }))} + /> + +
+ +
+ + +
+
+
+ ) : ( +
+
+ +
+ {loading ? ( +
+ +
+ ) : servers.length === 0 ? ( +
+ No MCP servers configured yet. +
+ ) : ( +
    + {servers.map((server) => ( +
  • +
    +
    + {server.display_name} +
    +
    + {server.url} +
    +
    +
    + toggleEnabled(server, next)} + aria-label="Enable server" + /> + + + +
    +
    • + ))} +
+ )} +
+ )} + + + ); +} diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index ce02b0da18..c4a5f34273 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -57,6 +57,7 @@ import { getProviderCapabilities, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, + providerSupportsBuiltinWebFetch, providerSupportsBuiltinWebSearch, } from "./provider-capabilities"; import { ChatRuntimeProvider } from "./runtime-provider"; @@ -71,6 +72,7 @@ import { CHAT_CODE_TOOLS_ENABLED_KEY, CHAT_IMAGE_TOOLS_ENABLED_KEY, CHAT_TOOLS_ENABLED_KEY, + CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, loadOptionalBool, useChatRuntimeStore, } from "./stores/chat-runtime-store"; @@ -726,7 +728,10 @@ export function ChatPage(): ReactElement { const reasoningCaps = getExternalReasoningCapabilities( provider?.providerType, selection.modelId, - { isReasoningProvider: provider?.isReasoningModel === true }, + { + isReasoningProvider: provider?.isReasoningModel === true, + baseUrl: provider?.baseUrl ?? null, + }, ); const state = useChatRuntimeStore.getState(); const preferredEffort = state.reasoningEffort; @@ -767,6 +772,8 @@ export function ChatPage(): ReactElement { : state.reasoningEffort; const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( provider?.providerType, + selection.modelId, + provider?.baseUrl, ); const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( provider?.providerType, @@ -779,6 +786,9 @@ export function ChatPage(): ReactElement { selection.modelId, provider?.baseUrl, ); + const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( + provider?.providerType, + ); // Kimi's k2.6/k2.5 default to thinking enabled on the server side // (per https://platform.kimi.ai/docs/models). Mirror that default // in the UI so the Think pill comes up clicked when the user picks @@ -801,6 +811,9 @@ export function ChatPage(): ReactElement { const storedImageToolsEnabled = loadOptionalBool( CHAT_IMAGE_TOOLS_ENABLED_KEY, ); + const storedWebFetchToolsEnabled = loadOptionalBool( + CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, + ); const nextToolsEnabled = supportsBuiltinWebSearch ? isKimi ? false @@ -834,6 +847,7 @@ export function ChatPage(): ReactElement { supportsBuiltinWebSearch, supportsBuiltinCodeExecution, supportsBuiltinImageGeneration, + supportsBuiltinWebFetch, toolsEnabled: nextToolsEnabled, codeToolsEnabled: supportsBuiltinCodeExecution ? (storedCodeToolsEnabled ?? false) @@ -841,6 +855,10 @@ export function ChatPage(): ReactElement { imageToolsEnabled: supportsBuiltinImageGeneration ? (storedImageToolsEnabled ?? false) : false, + // Default Fetch off (Anthropic bills per fetch); deliberate opt-in. + webFetchToolsEnabled: supportsBuiltinWebFetch + ? (storedWebFetchToolsEnabled ?? false) + : false, }); }, [externalProvidersForChat, inferenceParams.checkpoint]); const canCompare = useMemo(() => { @@ -955,6 +973,7 @@ export function ChatPage(): ReactElement { { isReasoningProvider: selectedProvider?.isReasoningModel === true, + baseUrl: selectedProvider?.baseUrl ?? null, }, ); const preferredEffort = store.reasoningEffort; @@ -996,6 +1015,8 @@ export function ChatPage(): ReactElement { store.setCheckpoint(value, null); const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch( selectedProvider?.providerType, + selectedExternal?.modelId, + selectedProvider?.baseUrl, ); const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution( selectedProvider?.providerType, @@ -1008,6 +1029,9 @@ export function ChatPage(): ReactElement { selectedExternal?.modelId, selectedProvider?.baseUrl, ); + const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( + selectedProvider?.providerType, + ); // See sibling useEffect above: Kimi's k2.x default to thinking // enabled, so the Think pill comes up clicked. Search pill stays // off by default; mutual exclusion flips them via the composer. @@ -1026,6 +1050,9 @@ export function ChatPage(): ReactElement { const storedImageToolsEnabled = loadOptionalBool( CHAT_IMAGE_TOOLS_ENABLED_KEY, ); + const storedWebFetchToolsEnabled = loadOptionalBool( + CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, + ); const nextToolsEnabled = supportsBuiltinWebSearch ? isKimi ? false @@ -1037,6 +1064,10 @@ export function ChatPage(): ReactElement { ggufMaxContextLength: null, ggufNativeContextLength: null, activeNativePathToken: null, + // Clear previous-model counters; the relaxed external-provider + // render gate would otherwise show stale stats until the next + // completion overwrites them. + contextUsage: null, supportsReasoning: reasoningCaps.supportsReasoning, reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn, reasoningStyle: reasoningCaps.reasoningStyle, @@ -1063,6 +1094,7 @@ export function ChatPage(): ReactElement { supportsBuiltinWebSearch, supportsBuiltinCodeExecution, supportsBuiltinImageGeneration, + supportsBuiltinWebFetch, toolsEnabled: nextToolsEnabled, codeToolsEnabled: supportsBuiltinCodeExecution ? (storedCodeToolsEnabled ?? false) @@ -1070,6 +1102,9 @@ export function ChatPage(): ReactElement { imageToolsEnabled: supportsBuiltinImageGeneration ? (storedImageToolsEnabled ?? false) : false, + webFetchToolsEnabled: supportsBuiltinWebFetch + ? (storedWebFetchToolsEnabled ?? false) + : false, ...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }), }); return; @@ -1161,7 +1196,9 @@ export function ChatPage(): ReactElement { if (!saved) return; viewBeforeCompareRef.current = null; navigate({ to: "/chat", search: saved }); - // Restore context usage from the active thread's last assistant message. + // Restore usage from the last assistant message, but only if it + // matches the currently active checkpoint. Without this guard the + // relaxed render gate would show stale stats from another model. const threadId = saved.thread ?? useChatRuntimeStore.getState().activeThreadId; if (threadId) { @@ -1175,7 +1212,29 @@ export function ChatPage(): ReactElement { const usage = metadata?.contextUsage as ReturnType< typeof useChatRuntimeStore.getState >["contextUsage"]; - if (usage) useChatRuntimeStore.getState().setContextUsage(usage); + if (!usage) return; + const store = useChatRuntimeStore.getState(); + const activeCheckpoint = store.params.checkpoint; + const usageModelId = + (usage as { modelId?: unknown }).modelId; + // Scope by modelId when present; reject if no active checkpoint + // (model-scoped usage cannot be attributed to "nothing"). + if (typeof usageModelId === "string" && usageModelId) { + if (!activeCheckpoint || usageModelId !== activeCheckpoint) { + return; + } + } + // For local turns, also require the restored count to fit in + // the active window. Skip when unknown (external provider). + const limit = store.ggufContextLength; + if ( + typeof limit === "number" && + limit > 0 && + (usage.totalTokens ?? 0) > limit + ) { + return; + } + store.setContextUsage(usage); }) .catch((error) => { if (!isExpectedBackgroundChatStorageError(error)) { @@ -1491,11 +1550,13 @@ export function ChatPage(): ReactElement { ) : null}
- {view.mode === "single" && ggufContextLength && contextUsage ? ( + {view.mode === "single" && contextUsage ? ( s.activeThreadId); const openAiApiKeyForSection = activeExternalProvider ? getExternalProviderApiKey(activeExternalProvider.id) || null : null; @@ -1305,6 +1316,28 @@ export function ChatSettingsPanel({
) : null} + {showFastModeControl ? ( +
+
+ + Fast mode + + + Beta. Up to 2.5x higher output tokens per second on + Claude Opus 4.6 and 4.7 at 6x standard Opus pricing. + Switching between fast and standard invalidates the + prompt cache and is incompatible with the Priority + service tier. + +
+ +
+ ) : null} ) : null} @@ -1841,6 +1874,12 @@ export function ChatSettingsPanel({
)} + + {!isExternalModel ? ( + + + + ) : null} s.mcpEnabledForChat); + const setMcpEnabledForChat = useChatRuntimeStore( + (s) => s.setMcpEnabledForChat, + ); + const [enabledServerCount, setEnabledServerCount] = useState( + null, + ); + const [dialogOpen, setDialogOpen] = useState(false); + const [refreshTick, setRefreshTick] = useState(0); + + useEffect(() => { + let cancelled = false; + listMcpServers() + .then((rows) => { + if (cancelled) return; + setEnabledServerCount(rows.filter((row) => row.is_enabled).length); + }) + .catch(() => { + if (!cancelled) setEnabledServerCount(0); + }); + return () => { + cancelled = true; + }; + }, [refreshTick]); + + return ( +
+
+
+ + Use MCP Servers + + + When on, every server marked enabled in the manage dialog is + attached to this chat's tool list. + +
+ +
+
+ + {enabledServerCount === null + ? "Loading…" + : enabledServerCount === 0 + ? "No servers configured" + : `${enabledServerCount} server${enabledServerCount === 1 ? "" : "s"} enabled`} + + +
+ { + setDialogOpen(next); + if (!next) setRefreshTick((tick) => tick + 1); + }} + /> +
+ ); +} + function ChatTemplateFields() { const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate); const override = useChatRuntimeStore((s) => s.chatTemplateOverride); diff --git a/studio/frontend/src/features/chat/components/context-usage-bar.tsx b/studio/frontend/src/features/chat/components/context-usage-bar.tsx index 91bdd35caa..e15a3e0b4f 100644 --- a/studio/frontend/src/features/chat/components/context-usage-bar.tsx +++ b/studio/frontend/src/features/chat/components/context-usage-bar.tsx @@ -28,37 +28,66 @@ function getSeverityColor(percent: number): { export const ContextUsageBar: FC<{ used: number; - total: number; + // null on external providers (no known window); bar then hides the ratio. + total?: number | null; cached?: number; + // Anthropic-only (billed at the write premium). + cacheWrites?: number; promptTokens?: number; completionTokens?: number; className?: string; -}> = ({ used, total, cached, promptTokens, completionTokens, className }) => { - if (total <= 0) return null; +}> = ({ + used, + total, + cached, + cacheWrites, + promptTokens, + completionTokens, + className, +}) => { + const hasKnownLimit = typeof total === "number" && total > 0; + const hasUsageDetails = + promptTokens !== undefined || + completionTokens !== undefined || + (cached !== undefined && cached > 0) || + (cacheWrites !== undefined && cacheWrites > 0); - const percent = Math.min((used / total) * 100, 100); - const severity = getSeverityColor(percent); + // Nothing to show: no limit and no per-turn counters. + if (!hasKnownLimit && used <= 0 && !hasUsageDetails) return null; + + const percent = hasKnownLimit + ? Math.min((used / (total as number)) * 100, 100) + : null; + const severity = getSeverityColor(percent ?? 0); return (
-
- Context usage - - {percent.toFixed(1)}% - -
+ {hasKnownLimit && percent !== null ? ( +
+ Context usage + + {percent.toFixed(1)}% + +
+ ) : null} {promptTokens !== undefined && (
Prompt tokens @@ -98,20 +129,32 @@ export const ContextUsageBar: FC<{
)} + {cacheWrites !== undefined && cacheWrites > 0 && ( +
+ Cache writes + + {formatTokenCountFull(cacheWrites)} + +
+ )}
- Total + + {hasKnownLimit ? "Total" : "Total tokens"} + - {formatTokenCountFull(used)} / {formatTokenCountFull(total)} + {hasKnownLimit + ? `${formatTokenCountFull(used)} / ${formatTokenCountFull(total as number)}` + : formatTokenCountFull(used)}
- {percent > 85 && ( + {hasKnownLimit && percent !== null && percent > 85 ? (
Close to the context limit. Generation will stop at 100%. Increase Context Length in the chat Settings panel to keep going.
- )} + ) : null}
diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index dddf529068..b10253c280 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -37,6 +37,13 @@ export interface ExternalProviderConfig { updatedAt: number; } +// Gemini supports prompt caching, but the wire flow requires a +// separate POST to /v1beta/cachedContents to create the cache before +// the generateContent call can reference it; the boolean Studio +// currently emits on enable_prompt_caching is not enough on its own. +// Until that two-step orchestration ships we keep the picker off so +// the toggle does not silently no-op for Gemini users. See +// https://ai.google.dev/gemini-api/docs/caching. const PROMPT_CACHING_PROVIDER_TYPES = new Set(["openai", "anthropic"]); export function supportsProviderPromptCaching( diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 4726b11fcf..883dea3f3a 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -2,6 +2,14 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export { ChatPage } from "./chat-page"; +export { + getInferenceStatus, + listGgufVariants, + listLocalModels, + loadModel, + type LocalModelInfo, +} from "./api/chat-api"; +export type { GgufVariantDetail } from "./types/api"; export { ChatSettingsPanel, defaultInferenceParams, diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts index ae8c1f41ae..4efbb74f11 100644 --- a/studio/frontend/src/features/chat/presets/preset-policy.ts +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -248,7 +248,7 @@ interface BackendInferenceDefaults { export interface BackendInferenceEnvelope { is_gguf?: boolean; context_length?: number | null; - inference?: BackendInferenceDefaults; + inference?: BackendInferenceDefaults | null; } export function mergeBackendRecommendedInference({ diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index da1d6e3431..f550279826 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -71,18 +71,95 @@ export function clampReasoningEffortToLevels( } /** - * Output-token cap for any external provider request. Picked to stay below the - * tightest declared limit across the providers we ship (Anthropic Claude Opus - * tops out at 128k, GPT-5.x ~128k, Gemini 2.5 ~65k, DeepSeek 8k) while staying - * well above what a typical chat reply needs. The local-model path is not - * subject to this — local backends honour whatever the loaded context allows. - * - * If a user's stored maxTokens (e.g. carried over from a prior local-model - * session with a 128k+ context) exceeds this, chat-adapter clamps the - * outbound request so the provider does not 400 on it. + * Fallback cap for unknown providers / models. Prefer + * `getExternalMaxOutputTokens(providerType, modelId)` for the real cap. */ export const EXTERNAL_MAX_OUTPUT_TOKENS = 32768; +/** + * Per-model max-output caps from each provider's docs: + * OpenAI: developers.openai.com/api/docs/models/gpt-5.5 + * Anthropic: platform.claude.com/docs/en/about-claude/models + * Gemini: ai.google.dev/gemini-api/docs/models/gemini-3.1-pro-preview + * DeepSeek: api-docs.deepseek.com/quick_start/pricing (V4 family) + * Local-model path is unaffected. + */ +const EXTERNAL_MAX_OUTPUT_TOKENS_BY_MODEL: Array<{ + providerType: string; + prefixes: readonly string[]; + cap: number; +}> = [ + // OpenAI + { providerType: "openai", prefixes: ["gpt-5.5-pro", "gpt-5.5"], cap: 128000 }, + { providerType: "openai", prefixes: ["gpt-5.4-pro", "gpt-5.4"], cap: 65536 }, + { providerType: "openai", prefixes: ["gpt-5.3"], cap: 16384 }, + // Anthropic + { + providerType: "anthropic", + prefixes: ["claude-opus-4-7"], + cap: 128000, + }, + { + providerType: "anthropic", + prefixes: [ + "claude-opus-4-6", + "claude-sonnet-4-6", + "claude-opus-4-5", + "claude-sonnet-4-5", + "claude-haiku-4-5", + ], + cap: 64000, + }, + // Gemini + { + providerType: "gemini", + prefixes: ["gemini-3", "gemini-pro", "gemini-flash"], + cap: 65536, + }, + // DeepSeek (V4: deepseek-chat / deepseek-reasoner alias V4-flash). + { providerType: "deepseek", prefixes: ["deepseek"], cap: 384000 }, +]; + +/** + * Documented per-model output cap; unknown ids fall back to + * `EXTERNAL_MAX_OUTPUT_TOKENS` (32k). OpenRouter ids are + * `provider/model`; the prefix is stripped before matching. + */ +export function getExternalMaxOutputTokens( + providerType: string | null | undefined, + modelId: string | null | undefined, +): number { + if (!providerType || !modelId) return EXTERNAL_MAX_OUTPUT_TOKENS; + const normalized = modelId.trim().toLowerCase(); + if (!normalized) return EXTERNAL_MAX_OUTPUT_TOKENS; + const stripped = + providerType === "openrouter" && normalized.includes("/") + ? normalized.split("/").slice(-1)[0] + : normalized; + const effectiveProvider = + providerType === "openrouter" + ? _inferProviderFromOpenrouterId(normalized) ?? providerType + : providerType; + for (const entry of EXTERNAL_MAX_OUTPUT_TOKENS_BY_MODEL) { + if (entry.providerType !== effectiveProvider) continue; + if (entry.prefixes.some((prefix) => stripped.startsWith(prefix))) { + return entry.cap; + } + } + return EXTERNAL_MAX_OUTPUT_TOKENS; +} + +function _inferProviderFromOpenrouterId( + normalizedId: string, +): string | null { + // Map OpenRouter `provider/model` prefix to our internal providerType. + if (normalizedId.startsWith("openai/")) return "openai"; + if (normalizedId.startsWith("anthropic/")) return "anthropic"; + if (normalizedId.startsWith("google/")) return "gemini"; + if (normalizedId.startsWith("deepseek/")) return "deepseek"; + return null; +} + /** * Whether the external provider offers a built-in web-search tool that the * model invokes server-side. When `true`, the chat composer's Search button @@ -112,7 +189,27 @@ export const EXTERNAL_MAX_OUTPUT_TOKENS = 32768; */ export function providerSupportsBuiltinWebSearch( providerType: string | null | undefined, + modelId?: string | null | undefined, + baseUrl?: string | null | undefined, ): boolean { + // Gemini ships grounded search via `tools: [{googleSearch: {}}]` on + // every chat-capable model. Most image-tier ids (`-image`, + // `nano-banana`) reject text-tool wiring because the + // responseModalities path is mutually exclusive with text tools, but + // Google explicitly documents Search grounding on the Gemini 3 image + // family (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + // nano-banana-pro). Allow Search on those; hide on older image ids. + // Custom Gemini OpenAI-compat proxies (non-Google bases) skip the + // native translator on the backend, so native tool envelopes never + // reach them -- hide the pill there. + if (providerType === "gemini") { + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + const normalized = modelId?.trim().toLowerCase() ?? ""; + if (normalized && isGeminiImageModel(normalized)) { + return geminiImageModelAllowsGoogleSearch(normalized); + } + return true; + } return ( providerType === "openai" || providerType === "anthropic" || @@ -123,11 +220,9 @@ export function providerSupportsBuiltinWebSearch( /** * Whether the external provider exposes a server-side web_fetch tool - * that retrieves a single URL (text or PDF) and emits a document block. - * Only Anthropic ships one today (`web_fetch_20250910`); the chat - * composer pairs it with the Search pill because the typical workflow - * is "search returns URLs, fetch reads them" and the UI doesn't (yet) - * expose web_fetch as an independent toggle. + * (single URL, text or PDF) emitting a document block. Anthropic-only + * today (`web_fetch_20250910` / `web_fetch_20260209`). Gates the + * composer's standalone Fetch pill, independent of Search. */ export function providerSupportsBuiltinWebFetch( providerType: string | null | undefined, @@ -135,6 +230,30 @@ export function providerSupportsBuiltinWebFetch( return providerType === "anthropic"; } +/** + * Whether the active provider + model supports Anthropic fast-mode + * (`speed: "fast"` + `fast-mode-2026-02-01` header). Opus 4.6 / 4.7 + * only per https://platform.claude.com/docs/en/build-with-claude/fast-mode. + * Backend silently drops on unsupported models as a second defence. + */ +const ANTHROPIC_FAST_MODE_MODEL_PREFIXES = [ + "claude-opus-4-7", + "claude-opus-4-6", +] as const; + +export function providerSupportsFastMode( + providerType: string | null | undefined, + modelId: string | null | undefined, +): boolean { + if (providerType !== "anthropic") return false; + if (!modelId) return false; + // Family boundary ("" or "-") required so IDs like "claude-opus-4-70" + // / "claude-opus-4-7b" do not match. + return ANTHROPIC_FAST_MODE_MODEL_PREFIXES.some( + (prefix) => modelId === prefix || modelId.startsWith(`${prefix}-`), + ); +} + /** * Whether the selected external provider/model exposes a server-side * code-execution tool. Two providers ship one today: @@ -185,15 +304,21 @@ const OPENAI_CODE_EXECUTION_MODEL_PREFIXES = [ /** * Strict check that a provider configuration points at OpenAI's - * managed cloud (api.openai.com), as opposed to a custom OpenAI-compat - * backend (ollama / llama.cpp / vLLM / generic "custom" preset). The - * shell tool ONLY exists on OpenAI cloud; sending it to anything else - * 400s the request. Mirror of the backend's - * `is_openai_cloud = "api.openai.com" in self.base_url` guard. + * managed cloud (api.openai.com) or Azure OpenAI Foundry + * (*.openai.azure.com), as opposed to a custom OpenAI-compat backend + * (ollama / llama.cpp / vLLM / generic "custom" preset). The shell and + * image-generation tools only exist on cloud backends; sending them to + * anything else 400s the request. Mirror of the backend's + * `_is_openai_family_cloud` host check. */ function isOpenAICloudBaseUrl(baseUrl: string | null | undefined): boolean { if (!baseUrl) return true; // No override → uses the default openai.com base. - return baseUrl.trim().toLowerCase().includes("api.openai.com"); + try { + const host = new URL(baseUrl).hostname.toLowerCase(); + return host === "api.openai.com" || host.endsWith(".openai.azure.com"); + } catch { + return false; + } } export function providerSupportsBuiltinCodeExecution( @@ -214,6 +339,20 @@ export function providerSupportsBuiltinCodeExecution( normalized.startsWith(prefix), ); } + if (providerType === "gemini") { + // Gemini's `tools: [{codeExecution: {}}]` is supported on every + // chat-capable model. Image-tier ids (`-image`, `nano-banana`) + // reject text-tool wiring because the inline-image path is + // mutually exclusive with codeExecution. Custom Gemini + // OpenAI-compat proxies skip the native translator on the + // backend, so native codeExecution envelopes do not reach them. + // Wire-up lives in `_stream_gemini` on the backend; output comes + // back inline as executableCode/codeExecutionResult parts. See + // https://ai.google.dev/gemini-api/docs/code-execution. + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + if (isGeminiImageModel(normalized)) return false; + return normalized.startsWith("gemini-"); + } return false; } @@ -246,12 +385,75 @@ export function providerSupportsBuiltinImageGeneration( modelId: string | null | undefined, baseUrl?: string | null, ): boolean { - if (providerType !== "openai") return false; - if (!isOpenAICloudBaseUrl(baseUrl)) return false; const normalized = modelId?.trim().toLowerCase() ?? ""; if (!normalized) return false; - return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => - normalized.startsWith(prefix), + if (providerType === "openai") { + if (!isOpenAICloudBaseUrl(baseUrl)) return false; + return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ); + } + if (providerType === "gemini") { + // Gemini's Nano Banana image-output ids carry either `-image` (e.g. + // `gemini-2.5-flash-image`, `gemini-3.1-flash-image-preview`) or the + // `nano-banana` alias (`nano-banana-pro-preview`). The backend flips + // generationConfig.responseModalities to ["TEXT", "IMAGE"] when one + // is picked, and translates inlineData parts into the same image_b64 + // tool_end envelope the OpenAI path emits so the chat UI renders the + // picture inline. Custom Gemini OpenAI-compat proxies skip the + // native translator on the backend, so hide the image pill there. + // See https://ai.google.dev/gemini-api/docs/image-generation. + if (isGeminiCustomOpenAICompatBase(baseUrl)) return false; + return normalized.includes("-image") || normalized.includes("nano-banana"); + } + return false; +} + +/** + * Whether `modelId` is a Gemini image-output id (Nano Banana family). + * Mirrors the backend's `is_image_picker_model` guard so the frontend + * hides text-only tool pills (web_search, code_execution) for these. + */ +function isGeminiImageModel(modelId: string): boolean { + const m = modelId.toLowerCase(); + return m.includes("-image") || m.includes("nano-banana"); +} + +/** + * Whether the saved Gemini connection points at a custom + * OpenAI-compatible gateway (any non-Google host). The backend + * `_is_openai_compatible` mirrors this to route those connections + * through `/chat/completions` instead of the native translator, so + * native Gemini tool envelopes (googleSearch, codeExecution, + * responseModalities) never reach them. Hide the corresponding + * Studio pills here so the request, builder, and UI agree. + */ +export function isGeminiCustomOpenAICompatBase( + baseUrl: string | null | undefined, +): boolean { + if (!baseUrl) return false; + try { + const host = new URL(baseUrl).hostname.toLowerCase(); + return host.length > 0 && host !== "generativelanguage.googleapis.com"; + } catch { + return false; + } +} + +/** + * Whether the given Gemini image model supports `tools: [{googleSearch: {}}]`. + * Google documents Search grounding on the Gemini 3 image family + * (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview, + * "Nano Banana Pro"); older image ids (gemini-2.5-flash-image) reject + * it with "Search as tool is not enabled for this model". + */ +function geminiImageModelAllowsGoogleSearch(modelId: string): boolean { + const m = modelId.toLowerCase(); + return ( + m.startsWith("gemini-3-pro-image") || + m.startsWith("gemini-3.1-flash-image") || + m.startsWith("nano-banana-pro") || + m.startsWith("nano-banana-2") ); } @@ -326,7 +528,20 @@ const PROVIDER_CAPABILITIES: Record = { presencePenalty: false, }, mistral: OPENAI_COMPAT_BASE, - gemini: OPENAI_COMPAT_BASE, + // Gemini's native generationConfig accepts temperature, topP, topK and + // presencePenalty (plus a separate frequencyPenalty we do not surface + // today). minP and repetitionPenalty are not part of the contract -- + // see https://ai.google.dev/api/rest/v1beta/GenerationConfig. Backend + // request shaping lives in _stream_gemini in + // studio/backend/core/inference/external_provider.py. + gemini: { + temperature: true, + topP: true, + topK: true, + minP: false, + repetitionPenalty: false, + presencePenalty: true, + }, // Kimi k2.5/k2.6 are reasoning-class — the API locks temperature and // top_p to fixed defaults and 400s on any other value: // "invalid temperature: only 1 is allowed for this model". @@ -538,6 +753,119 @@ function resolveKimiReasoningCapabilities(modelId: string): ExternalReasoningCap return withEnableThinkingStyle(); } +// Gemini's thinking ladder. +// - Gemini 3.x (3 / 3.1 / 3.5, Pro + Flash + Flash-Lite) and the +// gemini-pro-latest / gemini-flash-latest aliases use the new +// `thinkingConfig.thinkingLevel` string field (LOW/MEDIUM/HIGH/ +// MINIMAL). Pro tier rejects MINIMAL. +// - Gemini 2.5 Flash + 2.5 Pro stay on the integer +// `thinkingConfig.thinkingBudget` (0=off on Flash, -1=dynamic, +// N>0=cap; Pro rejects 0). +// - 2.5 Flash-Lite: no native thinking surfaced; leave it off. +// - Image-tier ids (`*-image*`, `nano-banana-pro-preview`): image +// generation path -- no reasoning controls. +const GEMINI3_PRO_PREFIXES = [ + "gemini-3.5-pro", + "gemini-3.1-pro", + "gemini-3-pro-preview", + "gemini-pro-latest", +]; +const GEMINI3_FLASH_PREFIXES = [ + "gemini-3.5-flash", + "gemini-3.1-flash", + "gemini-3-flash", + "gemini-flash-latest", + "gemini-flash-lite-latest", +]; +const GEMINI25_PRO_PREFIXES = [ + "gemini-2.5-pro", +]; +const GEMINI25_FLASH_PREFIXES = [ + "gemini-2.5-flash", +]; +const GEMINI_IMAGE_HINTS = [ + "-image", + "nano-banana", +]; +function resolveGeminiReasoningCapabilities( + modelId: string, +): ExternalReasoningCapabilities { + const m = modelId.toLowerCase(); + if (GEMINI_IMAGE_HINTS.some((h) => m.includes(h))) { + // Image generation; no thinking knob. + return withEnableThinkingStyle(); + } + // Gemini 2.5 Flash-Lite supports `thinkingBudget` with `0` = off and + // a positive range starting at 512 (the backend maps "minimal" to + // that floor at external_provider._stream_gemini). Check this branch + // BEFORE the broader `gemini-2.5-flash` prefix. + // https://ai.google.dev/gemini-api/docs/thinking + if (m.startsWith("gemini-2.5-flash-lite")) { + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + reasoningEffortLevels: [ + "none", + "minimal", + "low", + "medium", + "high", + "max", + ] as const, + }); + } + if (GEMINI3_PRO_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 3.x Pro: thinkingLevel supports low/medium/high per + // https://ai.google.dev/gemini-api/docs/thinking and + // https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro. + // Cannot fully disable thinking; "minimal" is rejected on Pro. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: ["low", "medium", "high"] as const, + }); + } + if (GEMINI3_FLASH_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 3 Flash: thinkingLevel minimal/low/medium/high. Minimal + // is the closest to "off" Google offers on Gemini 3. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: [ + "minimal", + "low", + "medium", + "high", + ] as const, + }); + } + if (GEMINI25_PRO_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 2.5 Pro: thinkingBudget cannot be 0 (API rejects with + // "only works in thinking mode"); backend coerces to a small + // positive budget. The picker still hides the off switch. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: false, + reasoningEffortLevels: ["low", "medium", "high", "max"] as const, + }); + } + if (GEMINI25_FLASH_PREFIXES.some((p) => m.startsWith(p))) { + // Gemini 2.5 Flash: thinkingBudget supports 0 = off cleanly. + return withReasoningEffortStyle({ + supportsReasoning: true, + supportsReasoningOff: true, + reasoningEffortLevels: [ + "none", + "low", + "medium", + "high", + "max", + ] as const, + }); + } + return withEnableThinkingStyle(); +} + function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoningCapabilities { if (modelId === "magistral-medium-latest") { return withReasoningEffortStyle({ @@ -560,6 +888,8 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning export interface ExternalReasoningResolveOptions { /** vLLM connection flagged as a reasoning model in provider config. */ isReasoningProvider?: boolean; + /** Provider base URL; used to detect custom Gemini OAI-compat gateways. */ + baseUrl?: string | null; } // vLLM has no per-model reasoning signal on OpenAI-compat — pin via user toggle. @@ -635,6 +965,16 @@ export function getExternalReasoningCapabilities( } if (isKimiProvider) return resolveKimiReasoningCapabilities(modelForMatching); if (isMistralProvider) return resolveMistralReasoningCapabilities(modelForMatching); + if (normalizedProvider === "gemini") { + // Custom Gemini OAI-compat gateways (LiteLLM, proxies) route + // through /chat/completions which drops the Gemini-native + // thinkingConfig payload. Hide the native thinking ladder so the + // UI does not advertise a control the backend cannot honor. + if (isGeminiCustomOpenAICompatBase(options?.baseUrl)) { + return withEnableThinkingStyle(); + } + return resolveGeminiReasoningCapabilities(modelForMatching); + } if (!isOpenAIProvider && !isAnthropicProvider) { return withEnableThinkingStyle(); } diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index d01383b309..21be5f6e3e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -826,17 +826,24 @@ function useStudioRuntimeAdapters(): StudioRuntimeAdapters { completionTokens: number; totalTokens: number; cachedTokens: number; + cacheWriteTokens?: number; modelId?: string; } | undefined; const store = useChatRuntimeStore.getState(); - if ( - savedUsage && - store.ggufContextLength && - savedUsage.totalTokens <= store.ggufContextLength && - (!savedUsage.modelId || - savedUsage.modelId === store.params.checkpoint) - ) { + // Window check applies only when a local GGUF window is known; + // external providers have ggufContextLength === null. + const withinLocalLimit = + !store.ggufContextLength || + (savedUsage?.totalTokens ?? 0) <= store.ggufContextLength; + // Legacy unscoped usage (no modelId) is only trusted when a + // known local window bounds the totals, so we can't misattribute + // an old local turn to a newly-selected external provider. + const modelMatches = savedUsage?.modelId + ? savedUsage.modelId === store.params.checkpoint + : typeof store.ggufContextLength === "number" && + store.ggufContextLength > 0; + if (savedUsage && withinLocalLimit && modelMatches) { store.setContextUsage(savedUsage); } diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 8292b1e5d2..c27b40da7b 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -21,9 +21,24 @@ import { isTauri } from "@/lib/api-base"; import { isMultimodalResponse } from "./types/api"; import { getImageInputUnavailableReason } from "./utils/image-input-support"; import { useAui } from "@assistant-ui/react"; -import { ArrowUpIcon, BookOpenIcon, FileTextIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; -import { useRagStore } from "@/features/rag/stores/rag-store"; +import { + ArrowUpIcon, + BookOpenIcon, + DownloadIcon, + FileTextIcon, + GlobeIcon, + HeadphonesIcon, + LightbulbIcon, + LightbulbOffIcon, + MicIcon, + PlusIcon, + SquareIcon, + XIcon, +} from "lucide-react"; +import { Image03Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { subscribeToJobEvents } from "@/features/rag/api/rag-api"; +import { useRagStore } from "@/features/rag/stores/rag-store"; import { toast } from "@/lib/toast"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers"; @@ -36,6 +51,7 @@ import { getExternalReasoningCapabilities, providerSupportsBuiltinCodeExecution, providerSupportsBuiltinImageGeneration, + providerSupportsBuiltinWebFetch, } from "./provider-capabilities"; import { type CompositionEvent, @@ -371,6 +387,12 @@ export function SharedComposer({ ); const ragToolEnabled = useChatRuntimeStore((s) => s.ragToolEnabled); const setRagToolEnabled = useChatRuntimeStore((s) => s.setRagToolEnabled); + const webFetchToolsEnabled = useChatRuntimeStore( + (s) => s.webFetchToolsEnabled, + ); + const setWebFetchToolsEnabled = useChatRuntimeStore( + (s) => s.setWebFetchToolsEnabled, + ); const lastOpenRouterChosenModel = useChatRuntimeStore( (s) => s.lastOpenRouterChosenModel, ); @@ -411,6 +433,7 @@ export function SharedComposer({ { isReasoningProvider: selectedExternalProvider?.isReasoningModel === true, + baseUrl: selectedExternalProvider?.baseUrl ?? null, }, ) : null; @@ -461,17 +484,43 @@ export function SharedComposer({ effectiveExternalModelId, selectedExternalProvider?.baseUrl, ); - const searchDisabled = - !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); - const codeDisabled = - !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution); - // Images pill is only ever lit on OpenAI cloud's Responses-API models. - // No local tool runtime fallback because the only image-generation - // server tool we wire today is OpenAI's; local models cannot dispatch - // it. Hidden entirely when the active model does not advertise it so - // the pill row stays compact for providers without the capability. + const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch( + selectedExternalProvider?.providerType, + ); + // Gemini rejects codeExecution alongside image modalities. Search is + // blocked on older Gemini image ids but allowed on Gemini 3 image + // models -- supportsBuiltinWebSearch already encodes the per-model + // allowance, so we only disable Code unconditionally in Gemini + // image mode. + const isExternalGemini = selectedExternalProvider?.providerType === "gemini"; const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration; + const imageModeDisablesCode = + isExternalGemini && imageToolsEnabled && !imageDisabled; + // Image-tier Gemini models always reject codeExecution and reject + // web_search on older ids (Gemini 3.x Pro/Flash allow it -- encoded + // in supportsBuiltinWebSearch). Don't let the local `supportsTools` + // runtime flag re-enable a pill the Gemini backend will silently + // drop. Detect "external provider is Gemini AND model is image-tier" + // and gate strictly on the provider builtin support. + const isGeminiImageTier = + isExternalGemini && supportsBuiltinImageGeneration; + const searchDisabled = + !modelLoaded || + (isGeminiImageTier + ? !supportsBuiltinWebSearch + : !(supportsTools || supportsBuiltinWebSearch)); + const codeDisabled = + !modelLoaded || + (isGeminiImageTier + ? true + : !(supportsTools || supportsBuiltinCodeExecution)) || + imageModeDisablesCode; + // Images pill is only ever lit on OpenAI cloud's Responses-API models + // and Gemini Nano Banana family. No local tool runtime fallback. const showImagePill = supportsBuiltinImageGeneration; + // Fetch pill: Anthropic-only (web_fetch_20250910 / web_fetch_20260209). + const webFetchDisabled = !modelLoaded || !supportsBuiltinWebFetch; + const showWebFetchPill = supportsBuiltinWebFetch; // Backwards-compatible alias for any other call site that may still // reference `toolsDisabled` (rare; both pills used it before). const toolsDisabled = codeDisabled; @@ -1082,7 +1131,7 @@ export function SharedComposer({ side="bottom" variant="ghost" size="icon" - className="size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:border-muted-foreground/15 dark:hover:bg-muted-foreground/30" + className="size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:hover:bg-muted-foreground/30" onClick={() => { // The picker accepts both image and audio. Don't gate the // button on image-availability — addFiles still filters @@ -1317,7 +1366,11 @@ export function SharedComposer({ imageToolsEnabled ? "Disable image generation" : "Enable image generation" } > - + Images )} @@ -1344,6 +1397,23 @@ export function SharedComposer({ RAG + {showWebFetchPill && ( + + )}
{dictationSupported && ( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 1d1df695eb..85407201a4 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -14,7 +14,9 @@ import { DEFAULT_INFERENCE_PARAMS, type InferenceParams, } from "../types/runtime"; -import { isExternalModelId } from "../external-providers"; +import { isExternalModelId, parseExternalModelId } from "../external-providers"; +import { getExternalMaxOutputTokens } from "../provider-capabilities"; +import { useExternalProvidersStore } from "./external-providers-store"; import { loadChatSettingsWithLegacyImport, savePersistedChatSettingsPatch, @@ -27,6 +29,9 @@ export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled"; export const CHAT_RAG_TOOL_ENABLED_KEY = "unsloth_chat_rag_tool_enabled"; +export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled"; +export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = + "unsloth_chat_web_fetch_tools_enabled"; // External provider selection is encoded into `params.checkpoint` as // `external::::`. PersistedChatSettings deliberately @@ -64,6 +69,12 @@ function saveLastExternalCheckpoint(value: string | null): void { } export type ReasoningStyle = "enable_thinking" | "reasoning_effort"; +export type PendingImageEditReference = { + threadId: string | null; + openaiImageGenerationCallId: string; + openaiResponseId?: string; + openaiReasoningItem?: unknown; +}; export type ReasoningEffort = | "none" | "minimal" @@ -264,10 +275,23 @@ type ChatRuntimeStore = { * receive the tool because their runtime cannot dispatch it. */ supportsBuiltinImageGeneration: boolean; + /** + * Whether the active external provider exposes a server-side + * web_fetch tool (Anthropic's `web_fetch_20250910` / + * `web_fetch_20260209`). Gates the composer's Fetch pill, + * independent of Search. + */ + supportsBuiltinWebFetch: boolean; toolsEnabled: boolean; ragToolEnabled: boolean; codeToolsEnabled: boolean; imageToolsEnabled: boolean; + mcpEnabledForChat: boolean; + /** + * Fetch pill state, independent of `toolsEnabled` (Search). Only + * consulted when `providerSupportsBuiltinWebFetch` is true. + */ + webFetchToolsEnabled: boolean; toolStatus: string | null; generatingStatus: string | null; autoHealToolCalls: boolean; @@ -289,11 +313,14 @@ type ChatRuntimeStore = { settingsPanelOpen: boolean; pendingAudioBase64: string | null; pendingAudioName: string | null; + pendingImageEditReference: PendingImageEditReference | null; contextUsage: { promptTokens: number; completionTokens: number; totalTokens: number; cachedTokens: number; + // Anthropic-only; optional so pre-cache-stats persisted entries load. + cacheWriteTokens?: number; } | null; modelLoading: boolean; activeNativePathToken: string | null; @@ -333,6 +360,8 @@ type ChatRuntimeStore = { setToolsEnabled: (enabled: boolean, options?: { persist?: boolean }) => void; setCodeToolsEnabled: (enabled: boolean) => void; setImageToolsEnabled: (enabled: boolean) => void; + setMcpEnabledForChat: (enabled: boolean) => void; + setWebFetchToolsEnabled: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; @@ -345,6 +374,10 @@ type ChatRuntimeStore = { setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; + setPendingImageEditReference: ( + reference: PendingImageEditReference | null, + ) => void; + clearPendingImageEditReference: () => void; setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void; setRagSource: (source: RagSource) => void; setRagMode: (mode: RagMode) => void; @@ -397,6 +430,7 @@ const PERSISTED_INFERENCE_PARAM_KEYS = [ "maxTokens", "systemPrompt", "trustRemoteCode", + "fastMode", ] as const satisfies readonly PersistedInferenceParamKey[]; const SCALAR_SETTING_KEYS = [ @@ -589,11 +623,14 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsBuiltinWebSearch: false, supportsBuiltinCodeExecution: false, supportsBuiltinImageGeneration: false, + supportsBuiltinWebFetch: false, toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false), // Defaults off; hydratePersistedSettings nudges it on for existing users. ragToolEnabled: loadBool(CHAT_RAG_TOOL_ENABLED_KEY, false), codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false), imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), + mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false), + webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false), toolStatus: null, generatingStatus: null, autoHealToolCalls: true, @@ -614,6 +651,7 @@ export const useChatRuntimeStore = create((set, get) => ({ settingsPanelOpen: false, pendingAudioBase64: null, pendingAudioName: null, + pendingImageEditReference: null, contextUsage: null, modelLoading: false, activeNativePathToken: null, @@ -683,7 +721,14 @@ export const useChatRuntimeStore = create((set, get) => ({ if (state.settingsHydrated && hasKeys(changedParams)) { saveSettingsPatch({ inferenceParams: changedParams }); } - return { params }; + // Mirror setCheckpoint: the local model load path can mutate + // params.checkpoint via setParams() before setCheckpoint runs, + // leaving stale per-turn counters under the new checkpoint. + const checkpointChanged = state.params.checkpoint !== params.checkpoint; + return { + params, + ...(checkpointChanged ? { contextUsage: null } : {}), + }; }), setCustomPresets: (customPresets) => set(() => { @@ -747,12 +792,37 @@ export const useChatRuntimeStore = create((set, get) => ({ // mount, and a stale persisted local id would race against the // freshly-loaded model. See LAST_EXTERNAL_CHECKPOINT_KEY notes. saveLastExternalCheckpoint(isExternalModelId(modelId) ? modelId : null); + // Clear stale per-turn usage when the model changes; the relaxed + // external-provider render gate would otherwise show old counters + // until the next completion overwrites them. + const checkpointChanged = state.params.checkpoint !== modelId; + // Clamp maxTokens to the new model's cap on switch into an + // external model so a value carried over from a prior local + // session does not render above the slider's max. + let nextMaxTokens = state.params.maxTokens; + if (checkpointChanged && isExternalModelId(modelId)) { + const parsed = parseExternalModelId(modelId); + const provider = parsed + ? useExternalProvidersStore + .getState() + .providers.find((p) => p.id === parsed.providerId) + : null; + const cap = getExternalMaxOutputTokens( + provider?.providerType, + parsed?.modelId, + ); + if (nextMaxTokens > cap) { + nextMaxTokens = cap; + } + } return { params: { ...state.params, checkpoint: modelId, + maxTokens: nextMaxTokens, }, activeGgufVariant: ggufVariant ?? null, + ...(checkpointChanged ? { contextUsage: null } : {}), }; }), setActiveThreadId: (activeThreadId) => @@ -787,9 +857,11 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsBuiltinWebSearch: false, supportsBuiltinCodeExecution: false, supportsBuiltinImageGeneration: false, + supportsBuiltinWebFetch: false, toolsEnabled: false, codeToolsEnabled: false, imageToolsEnabled: false, + webFetchToolsEnabled: false, toolStatus: null, kvCacheDtype: null, loadedKvCacheDtype: null, @@ -802,6 +874,7 @@ export const useChatRuntimeStore = create((set, get) => ({ defaultChatTemplate: null, chatTemplateOverride: null, loadedChatTemplateOverride: null, + pendingImageEditReference: null, })); }, setReasoningEnabled: (reasoningEnabled, options) => @@ -891,6 +964,16 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, imageToolsEnabled); return { imageToolsEnabled }; }), + setMcpEnabledForChat: (mcpEnabledForChat) => + set(() => { + saveBool(CHAT_MCP_ENABLED_KEY, mcpEnabledForChat); + return { mcpEnabledForChat }; + }), + setWebFetchToolsEnabled: (webFetchToolsEnabled) => + set(() => { + saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled); + return { webFetchToolsEnabled }; + }), setToolStatus: (toolStatus) => set({ toolStatus }), setGeneratingStatus: (generatingStatus) => set({ generatingStatus }), setAutoHealToolCalls: (autoHealToolCalls) => @@ -930,5 +1013,9 @@ export const useChatRuntimeStore = create((set, get) => ({ set({ pendingAudioBase64: base64, pendingAudioName: name }), clearPendingAudio: () => set({ pendingAudioBase64: null, pendingAudioName: null }), + setPendingImageEditReference: (pendingImageEditReference) => + set({ pendingImageEditReference }), + clearPendingImageEditReference: () => + set({ pendingImageEditReference: null }), setContextUsage: (contextUsage) => set({ contextUsage }), })); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 1e6bcf8b87..d313b43438 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -143,6 +143,7 @@ export interface UnloadModelRequest { export interface InferenceStatusResponse { active_model: string | null; + model_identifier?: string | null; is_vision: boolean; is_gguf?: boolean; gguf_variant?: string | null; @@ -158,7 +159,7 @@ export interface InferenceStatusResponse { min_p?: number; presence_penalty?: number; trust_remote_code?: boolean; - }; + } | null; requires_trust_remote_code?: boolean; supports_reasoning?: boolean; reasoning_style?: "enable_thinking" | "reasoning_effort"; @@ -192,16 +193,60 @@ export interface AudioGenerationResponse { }>; } -export type OpenAIMessageContent = - | string - | Array< - | { type: "text"; text: string } - | { type: "image_url"; image_url: { url: string } } - >; +export type OpenAIReasoningSummaryPart = { + type: "summary_text"; + text: string; +}; + +export type OpenAIReasoningContentPart = { + type: "reasoning"; + id: string; + summary: OpenAIReasoningSummaryPart[]; + status?: "in_progress" | "completed" | "incomplete"; +}; + +export type OpenAIImageGenerationCallContentPart = { + type: "image_generation_call"; + id: string; + response_id?: string; +}; + +export type OpenAIMessageContentPart = + | { type: "text"; text: string } + | { type: "image_url"; image_url: { url: string } } + | OpenAIReasoningContentPart + | OpenAIImageGenerationCallContentPart; + +export type OpenAIMessageContent = string | OpenAIMessageContentPart[]; + +/** + * OpenAI Chat Completions tool_call shape. Assistant turns echo back + * function/tool calls as `tool_calls`; the matching tool result rides + * on a separate `role="tool"` message keyed by `tool_call_id`. + * `extra_content.google.thought_signature` is the Gemini-specific + * round-trip field the backend translator both emits (on `delta. + * tool_calls`) and consumes (when rebuilding the native functionCall + * part on the next turn). + */ +export interface OpenAIToolCallPart { + id?: string; + type?: "function"; + function?: { + name?: string; + arguments?: string; + }; + extra_content?: unknown; +} export interface OpenAIChatMessage { - role: "system" | "user" | "assistant"; - content: OpenAIMessageContent; + role: "system" | "user" | "assistant" | "tool"; + content: OpenAIMessageContent | null; + /** Assistant tool-call deltas, when the turn invoked a function tool. */ + tool_calls?: OpenAIToolCallPart[]; + /** `role="tool"` only: id matching `assistant.tool_calls[].id`. */ + tool_call_id?: string; + /** `role="tool"` only: name of the function that produced the result. */ + name?: string; } export interface OpenAIChatCompletionsRequest { @@ -242,7 +287,14 @@ export interface OpenAIChatCompletionsRequest { external_model?: string; encrypted_api_key?: string; provider_base_url?: string | null; - enable_prompt_caching?: boolean | null; + /** + * Boolean toggle for OpenAI/Anthropic ephemeral cache_control. For + * Gemini the backend also accepts the cached-content resource name + * (`cachedContents/...`) as a string, which is forwarded as + * `generationConfig.cachedContent` on the native streamGenerateContent + * request. + */ + enable_prompt_caching?: boolean | string | null; /** * OpenAI shell-tool container id captured from the prior response in * this chat thread. When set and the Code pill is on, the backend @@ -262,11 +314,30 @@ export interface OpenAIChatCompletionsRequest { * the Anthropic provider with `code_execution` in `enabled_tools`. */ anthropic_code_exec_container_id?: string | null; + /** + * Anthropic fast-mode toggle. Opus 4.6 / 4.7 only; backend drops + * silently on every other model + provider. See + * https://platform.claude.com/docs/en/build-with-claude/fast-mode + */ + fast_mode?: boolean | null; } export interface OpenAIChatDelta { role?: string; - content?: string; + content?: string | null; + /** + * Streamed assistant tool calls. The Gemini and OpenAI Responses + * translators emit incremental `tool_calls` deltas (function name + + * arguments fragments) so the chat-adapter can render tool cards as + * they arrive. + */ + tool_calls?: OpenAIToolCallPart[]; + /** + * Provider-specific passthrough. Gemini ships `thoughtSignature`, + * citations, `native_part`, etc., here so the round-trip can replay + * them on follow-up turns without bleeding into other providers. + */ + extra_content?: Record; } export interface OpenAIChatChunkChoice { diff --git a/studio/frontend/src/features/chat/types/runtime.ts b/studio/frontend/src/features/chat/types/runtime.ts index 2967584653..4c44ee1e9c 100644 --- a/studio/frontend/src/features/chat/types/runtime.ts +++ b/studio/frontend/src/features/chat/types/runtime.ts @@ -14,6 +14,12 @@ export interface InferenceParams { checkpoint: string; /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trustRemoteCode?: boolean; + /** + * Anthropic fast-mode toggle. Opus 4.6 / 4.7 only; higher OTPS at + * 6x standard Opus pricing. Default false. + * https://platform.claude.com/docs/en/build-with-claude/fast-mode + */ + fastMode?: boolean; } export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { @@ -28,6 +34,7 @@ export const DEFAULT_INFERENCE_PARAMS: InferenceParams = { systemPrompt: "", checkpoint: "", trustRemoteCode: false, + fastMode: false, }; export interface ChatModelSummary { diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index e07e1ddb1d..4e93a20bff 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -140,6 +140,11 @@ function sanitizeInferenceParams( if (typeof value.trustRemoteCode === "boolean") { params.trustRemoteCode = value.trustRemoteCode; } + // Mirror trustRemoteCode handling so the toggle survives reload + // and the /api/chat/settings round-trip. + if (typeof value.fastMode === "boolean") { + params.fastMode = value.fastMode; + } return hasKeys(params) ? params : undefined; } diff --git a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx index f30af7fcf9..ed590f226e 100644 --- a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx +++ b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx @@ -5,6 +5,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { getAuthToken } from "@/features/auth"; +import { useT } from "@/i18n"; import { toastError, toastSuccess } from "@/shared/toast"; import { Camera } from "lucide-react"; import { useMemo, useRef, useState } from "react"; @@ -37,6 +38,7 @@ function readPersistedProfile(): { displayName: string; avatarDataUrl: string | } export function ProfilePersonalizationPanel() { + const t = useT(); const displayName = useUserProfileStore((s) => s.displayName); const avatarDataUrl = useUserProfileStore((s) => s.avatarDataUrl); const setDisplayName = useUserProfileStore((s) => s.setDisplayName); @@ -60,11 +62,11 @@ export function ProfilePersonalizationPanel() { setDisplayName(trimmed); const persisted = readPersistedProfile(); if (persisted && persisted.displayName === trimmed) { - toastSuccess("Profile name saved"); + toastSuccess(t("settings.profile.nameSaved")); } else { toastError( - "Could not persist profile name", - "Name updated for this session, but may not persist after reload.", + t("settings.profile.namePersistErrorTitle"), + t("settings.profile.namePersistErrorDescription"), ); } } @@ -78,17 +80,18 @@ export function ProfilePersonalizationPanel() { setAvatarDataUrl(dataUrl); const persisted = readPersistedProfile(); if (persisted && persisted.avatarDataUrl === dataUrl) { - toastSuccess("Profile photo updated"); + toastSuccess(t("settings.profile.photoUpdated")); } else { toastError( - "Could not persist profile photo", - "Photo updated for this session, but may not persist after reload.", + t("settings.profile.photoPersistErrorTitle"), + t("settings.profile.photoPersistErrorDescription"), ); } } catch (e) { - const message = e instanceof Error ? e.message : "Could not use this image."; + const message = + e instanceof Error ? e.message : t("settings.profile.imageUseError"); setImageError(message); - toastError("Could not update profile photo", message); + toastError(t("settings.profile.photoUpdateErrorTitle"), message); } }; @@ -115,7 +118,7 @@ export function ProfilePersonalizationPanel() { type="button" onClick={() => fileInputRef.current?.click()} className="absolute right-0 bottom-0 -translate-x-[15.625%] -translate-y-[15.625%] flex size-8 items-center justify-center rounded-full border border-border bg-background text-foreground shadow-sm transition-colors hover:bg-muted focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background" - aria-label="Change profile picture" + aria-label={t("settings.profile.changePicture")} > @@ -123,7 +126,7 @@ export function ProfilePersonalizationPanel() {
diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx index 16e99f4fae..1d98d18b56 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx @@ -3,6 +3,7 @@ import { Input } from "@/components/ui/input"; import type { ReactElement } from "react"; +import { LocalRecipeModelSelector } from "../../dialogs/models/local-recipe-model-selector"; import type { ModelConfig, ModelProviderConfig } from "../../types"; import { InlineField } from "./inline-field"; @@ -32,7 +33,9 @@ export function InlineModel(props: InlineModelProps): ReactElement { className="nodrag h-8 w-full text-xs" placeholder="https://api.example.com/v1" value={props.config.endpoint} - onChange={(event) => props.onUpdate({ endpoint: event.target.value })} + onChange={(event) => + props.onUpdate({ endpoint: event.target.value }) + } /> @@ -53,23 +56,32 @@ export function InlineModel(props: InlineModelProps): ReactElement { } // model_config branch - mirror the local-aware provider sync from the - // dialog path so inline edits do not leave stale "local" placeholders - // on external providers and fill the placeholder when switching to local. + // dialog path so inline edits clear stale local-only metadata without + // synthesizing the legacy "local" placeholder. const localNames = props.localProviderNames ?? new Set(); const modelConfig = props.config; - const handleProviderChange = (nextProvider: string) => { - const isLocal = localNames.has(nextProvider); - if (isLocal && !modelConfig.model.trim()) { - props.onUpdate({ provider: nextProvider, model: "local" }); - return; - } - if (!isLocal && modelConfig.model === "local") { - props.onUpdate({ provider: nextProvider, model: "" }); - return; - } - props.onUpdate({ provider: nextProvider }); - }; const isLinkedToLocal = localNames.has(modelConfig.provider); + const handleProviderChange = (nextProvider: string) => { + const nextIsLocal = localNames.has(nextProvider); + if (isLinkedToLocal !== nextIsLocal) { + props.onUpdate({ + provider: nextProvider, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }); + return; + } + props.onUpdate({ + provider: nextProvider, + ...(nextIsLocal + ? {} + : { + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }), + }); + }; return (
@@ -82,12 +94,38 @@ export function InlineModel(props: InlineModelProps): ReactElement { /> - props.onUpdate({ model: event.target.value })} - /> + {isLinkedToLocal ? ( + + props.onUpdate({ + model, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: variant ?? undefined, + }) + } + /> + ) : ( + + props.onUpdate({ + model: event.target.value, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }) + } + /> + )} void; + inputId?: string; + disabled?: boolean; + compact?: boolean; + className?: string; +}; + +function normalizeForSearch(value: string): string { + return value.toLowerCase().replace(/[\s_.-]/g, ""); +} + +function hasGgufSuffix(value: string | null | undefined): boolean { + return GGUF_SUFFIX_PATTERN.test(value ?? ""); +} + +function getModelLabel(model: LocalModelInfo): string { + return model.model_id?.trim() || model.display_name || model.id; +} + +function isDirectGguf(model: LocalModelInfo): boolean { + return model.path.toLowerCase().endsWith(".gguf"); +} + +function isExpandableGguf(model: LocalModelInfo): boolean { + return ( + !isDirectGguf(model) && + (hasGgufSuffix(model.id) || + hasGgufSuffix(model.display_name) || + hasGgufSuffix(model.model_id)) + ); +} + +function sourceLabel(model: LocalModelInfo): string { + switch (model.source) { + case "models_dir": + return "Models"; + case "hf_cache": + return "HF cache"; + case "lmstudio": + return "LM Studio"; + case "custom": + return "Custom folder"; + default: + return "Local"; + } +} + +type SelectedModelSummary = { + label: string; + source: string; + isGguf: boolean; +}; + +function getSelectedModelSummary( + value: string, + selectedModel: LocalModelInfo | null, + ggufVariant?: string | null, +): SelectedModelSummary { + if (!selectedModel) { + return { + label: value, + source: "Local model", + isGguf: Boolean(ggufVariant), + }; + } + + return { + label: getModelLabel(selectedModel), + source: sourceLabel(selectedModel), + isGguf: isDirectGguf(selectedModel) || isExpandableGguf(selectedModel), + }; +} + +function LocalGgufVariantList({ + repoId, + selectedVariant, + onSelect, +}: { + repoId: string; + selectedVariant?: string | null; + onSelect: (variant: string) => void; +}): ReactElement { + const [variants, setVariants] = useState(null); + const [defaultVariant, setDefaultVariant] = useState(null); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + + useEffect(() => { + let cancelled = false; + listGgufVariants(repoId) + .then((response) => { + if (cancelled) { + return; + } + setVariants(response.variants); + setDefaultVariant(response.default_variant); + }) + .catch((err) => { + if (cancelled) { + return; + } + setError( + err instanceof Error ? err.message : "Failed to load variants.", + ); + }) + .finally(() => { + if (!cancelled) { + setLoading(false); + } + }); + return () => { + cancelled = true; + }; + }, [repoId]); + + const sortedVariants = useMemo(() => { + if (!variants) { + return null; + } + return [...variants].sort((a, b) => { + if (a.quant === defaultVariant) { + return -1; + } + if (b.quant === defaultVariant) { + return 1; + } + if (a.downloaded !== b.downloaded) { + return a.downloaded ? -1 : 1; + } + return a.quant.localeCompare(b.quant); + }); + }, [defaultVariant, variants]); + + if (loading) { + return ( +
+ + Loading quantizations... +
+ ); + } + + if (error) { + return
{error}
; + } + + if (!sortedVariants || sortedVariants.length === 0) { + return ( +
+ No GGUF quantizations found for this model. +
+ ); + } + + return ( +
+
+ Quantization +
+
+ {sortedVariants.map((variant) => { + const selected = selectedVariant === variant.quant; + return ( + + ); + })} +
+
+ ); +} + +type SelectorTriggerProps = ComponentPropsWithoutRef<"button"> & { + value: string; + selectedModel: LocalModelInfo | null; + ggufVariant?: string | null; + inputId?: string; + disabled: boolean; + compact: boolean; + className?: string; +}; + +const SelectorTrigger = forwardRef( + function SelectorTrigger( + { + value, + selectedModel, + ggufVariant, + inputId, + disabled, + compact, + className, + ...triggerProps + }, + ref, + ): ReactElement { + const selected = getSelectedModelSummary(value, selectedModel, ggufVariant); + + return ( + + ); + }, +); + +function LocalModelRow({ + model, + selected, + expanded, + probing, + ggufVariant, + onSelectModel, + onSelectVariant, +}: { + model: LocalModelInfo; + selected: boolean; + expanded: boolean; + probing: boolean; + ggufVariant?: string | null; + onSelectModel: (model: LocalModelInfo) => void; + onSelectVariant: (modelId: string, variant: string) => void; +}): ReactElement { + const expandable = isExpandableGguf(model); + const directGguf = isDirectGguf(model); + + return ( +
+ + {expanded ? ( + onSelectVariant(model.id, variant)} + /> + ) : null} +
+ ); +} + +function LocalModelResults({ + loading, + error, + models, + value, + ggufVariant, + expandedModelId, + probingVariantModelId, + onRefresh, + onSelectModel, + onSelectVariant, +}: { + loading: boolean; + error: string | null; + models: LocalModelInfo[]; + value: string; + ggufVariant?: string | null; + expandedModelId: string | null; + probingVariantModelId: string | null; + onRefresh: () => void; + onSelectModel: (model: LocalModelInfo) => void; + onSelectVariant: (modelId: string, variant: string) => void; +}): ReactElement { + if (loading) { + return ( +
+ + Scanning local models... +
+ ); + } + + if (error) { + return ( +
+

{error}

+ +
+ ); + } + + if (models.length === 0) { + return ( +
+

No local models found.

+

+ Download a model or add a scan folder from Chat, then refresh this + list. +

+ + Open Chat model picker + +
+ ); + } + + return ( +
+ {models.map((model) => ( + + ))} +
+ ); +} + +export function LocalRecipeModelSelector({ + value, + ggufVariant, + onChange, + inputId, + disabled = false, + compact = false, + className, +}: LocalRecipeModelSelectorProps): ReactElement { + const [open, setOpen] = useState(false); + const [query, setQuery] = useState(""); + const [models, setModels] = useState([]); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [expandedModelId, setExpandedModelId] = useState(null); + const [probingVariantModelId, setProbingVariantModelId] = useState< + string | null + >(null); + const [refreshKey, setRefreshKey] = useState(0); + + const requestModelRefresh = useCallback(() => { + setLoading(true); + setError(null); + setRefreshKey((key) => key + 1); + }, []); + + const handleOpenChange = useCallback( + (nextOpen: boolean) => { + setOpen(nextOpen); + if (nextOpen) { + requestModelRefresh(); + } + }, + [requestModelRefresh], + ); + + useEffect(() => { + if (!open || refreshKey < 0) { + return; + } + let cancelled = false; + listLocalModels() + .then((response) => { + if (cancelled) { + return; + } + setModels(response.models); + }) + .catch((err) => { + if (cancelled) { + return; + } + setError( + err instanceof Error ? err.message : "Failed to list local models.", + ); + }) + .finally(() => { + if (!cancelled) { + setLoading(false); + } + }); + return () => { + cancelled = true; + }; + }, [open, refreshKey]); + + const selectedModel = useMemo( + () => models.find((model) => model.id === value) ?? null, + [models, value], + ); + + const filteredModels = useMemo(() => { + const needle = normalizeForSearch(query.trim()); + if (!needle) { + return models; + } + return models.filter((model) => { + const haystack = normalizeForSearch( + `${model.id} ${model.display_name} ${model.model_id ?? ""} ${model.path}`, + ); + return haystack.includes(needle); + }); + }, [models, query]); + + const selectModel = useCallback( + async (model: LocalModelInfo) => { + if (isExpandableGguf(model)) { + setExpandedModelId((current) => + current === model.id ? null : model.id, + ); + return; + } + if (!isDirectGguf(model)) { + setProbingVariantModelId(model.id); + try { + const response = await listGgufVariants(model.id); + if (response.variants.length > 0) { + setExpandedModelId(model.id); + return; + } + } catch { + // Non-GGUF local models commonly have no variant endpoint. Fall + // through to regular selection so users can still choose them. + } finally { + setProbingVariantModelId(null); + } + } + onChange(model.id, null); + setOpen(false); + }, + [onChange], + ); + + const selectVariant = useCallback( + (modelId: string, variant: string) => { + onChange(modelId, variant); + setOpen(false); + }, + [onChange], + ); + + return ( + + + + + +
+
+
+ setQuery(event.target.value)} + placeholder="Filter local models" + className="h-8 flex-1" + autoFocus={true} + /> + +
+
+ +
event.stopPropagation()} + > + +
+
+ + + ); +} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx index 368ae08acb..68f912bc57 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx @@ -1,12 +1,12 @@ // 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 { Checkbox } from "@/components/ui/checkbox"; import { Collapsible, CollapsibleContent, CollapsibleTrigger, } from "@/components/ui/collapsible"; -import { Checkbox } from "@/components/ui/checkbox"; import { Combobox, ComboboxContent, @@ -22,6 +22,7 @@ import type { ModelConfig } from "../../types"; import { CollapsibleSectionTriggerButton } from "../shared/collapsible-section-trigger"; import { FieldLabel } from "../shared/field-label"; import { NameField } from "../shared/name-field"; +import { LocalRecipeModelSelector } from "./local-recipe-model-selector"; type ModelConfigDialogProps = { config: ModelConfig; @@ -45,6 +46,7 @@ export function ModelConfigDialog({ const maxTokensId = `${config.id}-max-tokens`; const timeoutId = `${config.id}-timeout`; const extraBodyId = `${config.id}-inference-extra-body`; + const skipHealthCheckId = `${config.id}-skip-health-check`; const providerAnchorRef = useRef(null); const providerInputRef = useRef(config.provider); // Sync providerInputRef with the current provider value. Updating a ref in @@ -61,16 +63,25 @@ export function ModelConfigDialog({ onUpdate({ [key]: value } as Partial); }; - // Apply provider selection while keeping the local-provider model autofill - // consistent across both dropdown selection and free-typed + blur input. + // Apply provider selection while clearing model identifiers that only make + // sense for the previous provider locality. const applyProviderChange = (selectedProvider: string) => { - const isLocal = localProviderNames.has(selectedProvider); - if (isLocal && !config.model.trim()) { - onUpdate({ provider: selectedProvider, model: "local" }); + const nextIsLocal = localProviderNames.has(selectedProvider); + if (isLinkedToLocal !== nextIsLocal) { + onUpdate({ + provider: selectedProvider, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }); return; } - if (!isLocal && config.model === "local") { - onUpdate({ provider: selectedProvider, model: "" }); + if (!nextIsLocal) { + onUpdate({ + provider: selectedProvider, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }); return; } updateField("provider", selectedProvider); @@ -88,8 +99,8 @@ export function ModelConfigDialog({ Set up one reusable model choice for your AI steps

- Choose the provider connection, enter the exact model ID, then save any - generation defaults you want to reuse. + Choose the provider connection, enter the exact model ID, then save + any generation defaults you want to reuse.

@@ -144,15 +155,48 @@ export function ModelConfigDialog({ - updateField("model", event.target.value)} + hint={ + isLinkedToLocal + ? "Choose the local model Recipes should load before Run or Validate." + : "The exact model name sent to the connection." + } /> + {isLinkedToLocal ? ( + + onUpdate({ + model, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: variant ?? undefined, + }) + } + /> + ) : ( + + onUpdate({ + model: event.target.value, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }) + } + /> + )} + {isLinkedToLocal ? ( +

+ Recipes will load this model automatically. GGUF quantization is + saved with the preset. +

+ ) : null}
@@ -250,8 +294,12 @@ export function ModelConfigDialog({ } />
-
@@ -149,15 +187,28 @@ export function GithubCrawlerEasyView({
- handleModelChange(event.target.value)} - placeholder="unsloth/gemma-4-E2B-it-GGUF" - disabled={!modelConfig} + hint={ + isModelLinkedToLocal + ? "Choose the local model this recipe should load." + : "OpenAI-compatible model id." + } /> + {isModelLinkedToLocal ? ( + + ) : ( + handleModelChange(event.target.value)} + placeholder="unsloth/gemma-4-E2B-it-GGUF" + disabled={!modelConfig} + /> + )}
diff --git a/studio/frontend/src/features/recipe-studio/executions/tracker.ts b/studio/frontend/src/features/recipe-studio/executions/tracker.ts index d83f662fcf..97e70c66c4 100644 --- a/studio/frontend/src/features/recipe-studio/executions/tracker.ts +++ b/studio/frontend/src/features/recipe-studio/executions/tracker.ts @@ -40,6 +40,11 @@ type TrackRecipeExecutionParams = { onPreviewSuccess?: () => void; }; +export type TrackRecipeExecutionResult = { + success: boolean; + terminal: boolean; +}; + function isTerminalStatus(status: RecipeExecutionStatus): boolean { return status === "completed" || status === "error" || status === "cancelled"; } @@ -53,7 +58,8 @@ function normalizeCompletedProgress(input: { } { const { latestExecution, rows } = input; const progressTotal = - typeof latestExecution.progress?.total === "number" && latestExecution.progress.total > 0 + typeof latestExecution.progress?.total === "number" && + latestExecution.progress.total > 0 ? latestExecution.progress.total : latestExecution.rows > 0 ? latestExecution.rows @@ -92,7 +98,7 @@ export async function trackRecipeExecution({ onUpsert, onSetPreviewErrors, onPreviewSuccess, -}: TrackRecipeExecutionParams): Promise { +}: TrackRecipeExecutionParams): Promise { let done = false; let lastStatus: RecipeExecutionStatus = initialExecution.status; let completedEventPayload: Record | null = null; @@ -124,7 +130,9 @@ export async function trackRecipeExecution({ } const eventType = - typeof event.payload.type === "string" ? event.payload.type : event.event; + typeof event.payload.type === "string" + ? event.payload.type + : event.event; if (eventType === "job.started") { latestExecution = { @@ -163,7 +171,7 @@ export async function trackRecipeExecution({ error: typeof event.payload.error === "string" ? event.payload.error - : latestExecution.error ?? `${label} failed.`, + : (latestExecution.error ?? `${label} failed.`), }; onUpsert(latestExecution); return; @@ -178,6 +186,19 @@ export async function trackRecipeExecution({ return; } + if (eventType === "job.cancelled") { + lastStatus = "cancelled"; + done = true; + latestExecution = { + ...latestExecution, + status: "cancelled", + finishedAt: Date.now(), + error: latestExecution.error ?? "Run cancelled.", + }; + onUpsert(latestExecution); + return; + } + if (changed) { onUpsert(latestExecution); } @@ -189,6 +210,9 @@ export async function trackRecipeExecution({ try { while (!done) { const status = await getRecipeJobStatus(jobId); + if (done && isTerminalStatus(lastStatus)) { + break; + } const mappedStatus = mapJobStatus(status.status); lastStatus = mappedStatus; latestExecution = applyExecutionStatusSnapshot(latestExecution, status); @@ -200,18 +224,19 @@ export async function trackRecipeExecution({ } } } catch (error) { - const message = toErrorMessage(error, `${label} failed.`); - latestExecution = { - ...latestExecution, - status: "error", - error: message, - finishedAt: Date.now(), - }; - onUpsert(latestExecution); - if (notify) { - toastError(`${label} failed`, message); + const terminal = isTerminalStatus(lastStatus); + if (!terminal) { + const message = toErrorMessage(error, `${label} failed.`); + latestExecution = { + ...latestExecution, + error: message, + }; + onUpsert(latestExecution); + if (notify) { + toastError(`${label} failed`, message); + } + return { success: false, terminal: false }; } - return false; } finally { eventsAbortController.abort(); } @@ -220,7 +245,10 @@ export async function trackRecipeExecution({ for (let attempt = 0; attempt < 3; attempt += 1) { try { const finalStatus = await getRecipeJobStatus(jobId); - latestExecution = applyExecutionStatusSnapshot(latestExecution, finalStatus); + latestExecution = applyExecutionStatusSnapshot( + latestExecution, + finalStatus, + ); } catch { break; } @@ -229,19 +257,20 @@ export async function trackRecipeExecution({ } } - const eventAnalysis = completedEventPayload - ? completedEventPayload["analysis"] - : null; - const eventDataset = completedEventPayload - ? completedEventPayload["dataset"] - : null; + const completedPayload = completedEventPayload as Record< + string, + unknown + > | null; + const eventAnalysis = completedPayload ? completedPayload.analysis : null; + const eventDataset = completedPayload ? completedPayload.dataset : null; const eventProcessorArtifacts = - completedEventPayload && - typeof completedEventPayload["processor_artifacts"] === "object" && - completedEventPayload["processor_artifacts"] !== null - ? (completedEventPayload["processor_artifacts"] as Record) + completedPayload && + typeof completedPayload.processor_artifacts === "object" && + completedPayload.processor_artifacts !== null + ? (completedPayload.processor_artifacts as Record) : null; - const shouldFetchPreviewDataset = kind === "preview" && !Array.isArray(eventDataset); + const shouldFetchPreviewDataset = + kind === "preview" && !Array.isArray(eventDataset); const shouldFetchAnalysis = !completedEventPayload || typeof eventAnalysis !== "object" || @@ -262,9 +291,7 @@ export async function trackRecipeExecution({ ? normalizeAnalysis(analysisResult.value) : latestExecution.analysis; const datasetResponse = - datasetResult.status === "fulfilled" - ? datasetResult.value - : null; + datasetResult.status === "fulfilled" ? datasetResult.value : null; const dataset = datasetResponse ? normalizeDatasetRows(datasetResponse.dataset) : latestExecution.dataset; @@ -272,7 +299,10 @@ export async function trackRecipeExecution({ datasetResponse && typeof datasetResponse.total === "number" ? datasetResponse.total : latestExecution.datasetTotal; - const completedProgress = normalizeCompletedProgress({ latestExecution, rows }); + const completedProgress = normalizeCompletedProgress({ + latestExecution, + rows, + }); latestExecution = { ...latestExecution, @@ -285,7 +315,8 @@ export async function trackRecipeExecution({ datasetPage: 1, datasetPageSize: DATASET_PAGE_SIZE, error: null, - processor_artifacts: eventProcessorArtifacts ?? latestExecution.processor_artifacts, + processor_artifacts: + eventProcessorArtifacts ?? latestExecution.processor_artifacts, finishedAt: latestExecution.finishedAt ?? Date.now(), }; onUpsert(latestExecution); @@ -299,7 +330,7 @@ export async function trackRecipeExecution({ toastSuccess("Full run completed."); } } - return true; + return { success: true, terminal: true }; } if (lastStatus === "cancelled") { @@ -313,7 +344,7 @@ export async function trackRecipeExecution({ if (notify) { toastError(`${label} cancelled`, "The execution was cancelled."); } - return false; + return { success: false, terminal: true }; } latestExecution = { @@ -326,5 +357,5 @@ export async function trackRecipeExecution({ if (notify) { toastError(`${label} failed`, latestExecution.error ?? "Execution failed."); } - return false; + return { success: false, terminal: true }; } diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts index c3da5b1999..19a6a3004d 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts @@ -1,14 +1,11 @@ // 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 { useCallback, useEffect, useState } from "react"; -import { useShallow } from "zustand/react/shallow"; +import { getInferenceStatus, loadModel } from "@/features/chat"; import { toast } from "@/lib/toast"; import { toastError } from "@/shared/toast"; -import { - getInferenceStatus, - loadModel, -} from "@/features/chat/api/chat-api"; +import { useCallback, useEffect, useState } from "react"; +import { useShallow } from "zustand/react/shallow"; import { cancelRecipeJob, createRecipeJob, @@ -23,8 +20,8 @@ import type { import { DATASET_PAGE_SIZE, executionLabel, - normalizeRunName, normalizeDatasetRows, + normalizeRunName, toErrorMessage, withExecutionDefaults, } from "../executions/execution-helpers"; @@ -32,84 +29,243 @@ import { findResumableExecution, loadSortedRecipeExecutions, } from "../executions/hydration"; -import { createBaseExecutionRecord } from "../executions/runtime"; import { buildExecutionPayload, sanitizeExecutionRows, } from "../executions/run-settings"; +import { createBaseExecutionRecord } from "../executions/runtime"; import { trackRecipeExecution } from "../executions/tracker"; import { type RecipeRunSettings, useRecipeExecutionsStore, } from "../stores/recipe-executions"; -import type { RecipePayload, RecipePayloadResult } from "../utils/payload/types"; +import type { + RecipePayload, + RecipePayloadResult, +} from "../utils/payload/types"; -/** - * Auto-load the local model before running a recipe that uses it. - * - * Looks at payload.recipe.model_providers for any provider with is_local=true, - * finds the bound model_configs and asks the backend to load whichever model - * the first local-bound model_config points at. Skips when the inference - * server already has that exact model active. This removes the "open /chat - * first" prerequisite that users kept tripping on. - */ -async function ensureLocalModelLoaded( - payload: RecipePayload, -): Promise { +const GGUF_MODEL_PATTERN = /gguf/i; + +function collectUsedLlmModelAliases(payload: RecipePayload): Set { + const columns = Array.isArray(payload.recipe.columns) + ? payload.recipe.columns + : []; + const aliases = new Set(); + for (const column of columns) { + const columnType = column.column_type; + if (typeof columnType !== "string" || !columnType.startsWith("llm-")) { + continue; + } + const alias = column.model_alias; + if (typeof alias === "string" && alias.trim()) { + aliases.add(alias.trim()); + } + } + return aliases; +} + +type LocalModelSelection = { + target: string; + ggufVariant: string; + aliases: string[]; +}; + +type LocalModelLoadPlan = + | { selection: LocalModelSelection; error: null; legacyAliases?: never } + | { selection: null; error: string; legacyAliases?: never } + | { selection: null; error: null; legacyAliases: string[] }; + +type RestorableLocalModelSnapshot = { + selection: LocalModelSelection | null; + unrestorableLabel: string | null; +}; + +function getLocalProviderNames(payload: RecipePayload): Set { const providers = Array.isArray(payload.recipe.model_providers) - ? (payload.recipe.model_providers as Array>) + ? (payload.recipe.model_providers as Record[]) : []; const localProviderNames = new Set(); - for (const p of providers) { - if (p.is_local === true && typeof p.name === "string") { - localProviderNames.add(p.name); + for (const provider of providers) { + if (provider.is_local === true && typeof provider.name === "string") { + localProviderNames.add(provider.name); } } - if (localProviderNames.size === 0) { - return null; + return localProviderNames; +} + +function findUsedLocalModelConfigs( + payload: RecipePayload, + localProviderNames: Set, +): Record[] { + const usedAliases = collectUsedLlmModelAliases(payload); + if (usedAliases.size === 0) { + return []; } const modelConfigs = Array.isArray(payload.recipe.model_configs) - ? (payload.recipe.model_configs as Array>) + ? payload.recipe.model_configs : []; - const boundConfig = modelConfigs.find( - (c) => typeof c.provider === "string" && localProviderNames.has(c.provider), - ); + return modelConfigs.filter((config) => { + const provider = config.provider; + const alias = config.alias; + return ( + typeof provider === "string" && + localProviderNames.has(provider) && + typeof alias === "string" && + usedAliases.has(alias) + ); + }); +} + +function readLocalModelSelection( + boundConfig: Record, +): LocalModelLoadPlan { + const alias = + typeof boundConfig.alias === "string" ? boundConfig.alias : "local model"; const target = - typeof boundConfig?.model === "string" ? boundConfig.model.trim() : ""; + typeof boundConfig.model === "string" ? boundConfig.model.trim() : ""; + const ggufVariant = + typeof boundConfig.gguf_variant === "string" + ? boundConfig.gguf_variant.trim() + : ""; if (!target) { - return null; + return { + selection: null, + error: `Model config ${alias}: choose a local model before validating or running this recipe.`, + }; + } + if (target.toLowerCase() === "local") { + return { selection: null, error: null, legacyAliases: [alias] }; + } + return { selection: { target, ggufVariant, aliases: [alias] }, error: null }; +} + +function getLocalModelLoadPlan( + boundConfigs: Record[], +): LocalModelLoadPlan | null { + const selections = new Map(); + const legacyAliases: string[] = []; + for (const boundConfig of boundConfigs) { + const next = readLocalModelSelection(boundConfig); + if (next.error) { + return next; + } + if (next.legacyAliases) { + legacyAliases.push(...next.legacyAliases); + continue; + } + const selection = next.selection; + if (!selection) { + continue; + } + const key = `${selection.target.toLowerCase()}\u0000${selection.ggufVariant}`; + const existing = selections.get(key); + if (existing) { + existing.aliases.push(...selection.aliases); + continue; + } + selections.set(key, selection); } + if (legacyAliases.length > 0 && selections.size > 0) { + const aliases = [ + ...legacyAliases, + ...[...selections.values()].flatMap((selection) => selection.aliases), + ].join(", "); + return { + selection: null, + error: `Recipes found mixed legacy and selected local models. Reselect the same concrete local model for: ${aliases}.`, + }; + } + + if (legacyAliases.length > 0) { + return { selection: null, error: null, legacyAliases }; + } + + if (selections.size > 1) { + const aliases = [...selections.values()] + .flatMap((selection) => selection.aliases) + .join(", "); + return { + selection: null, + error: `Recipes supports one active local model per run. Select the same local model and GGUF variant for: ${aliases}.`, + }; + } + + const selection = [...selections.values()][0]; + return selection ? { selection, error: null } : null; +} + +function isDirectGgufTarget(target: string): boolean { + return target.toLowerCase().endsWith(".gguf"); +} + +function localSelectionMatchesActive(input: { + target: string; + ggufVariant: string; + activeModel: string | null | undefined; + activeVariant: string; +}): boolean { + const { target, ggufVariant, activeModel, activeVariant } = input; + if (!activeModel || activeModel.toLowerCase() !== target.toLowerCase()) { + return false; + } + return ( + activeVariant === ggufVariant || + (isDirectGgufTarget(target) && !ggufVariant) + ); +} + +async function isLocalModelAlreadyLoaded( + selection: LocalModelSelection, +): Promise { + const { target, ggufVariant } = selection; try { const status = await getInferenceStatus(); - if ( - status.active_model && - status.active_model.toLowerCase() === target.toLowerCase() - ) { - return null; - } + return localSelectionMatchesActive({ + target, + ggufVariant, + activeModel: status.model_identifier ?? status.active_model, + activeVariant: status.gguf_variant?.trim() ?? "", + }); } catch { // Fall through to load attempt; the backend will re-error if needed. + return false; } +} - const toastId = toast.loading(`Loading ${target}…`, { +async function loadLocalModelSelection( + selection: LocalModelSelection, +): Promise { + const { target, ggufVariant } = selection; + const modelLabel = ggufVariant ? `${target} (${ggufVariant})` : target; + const toastId = toast.loading(`Loading ${modelLabel}...`, { description: "Starting the local inference server for this recipe.", }); try { - const isGguf = /gguf/i.test(target); + const isGguf = GGUF_MODEL_PATTERN.test(target) || Boolean(ggufVariant); await loadModel({ + // biome-ignore lint/style/useNamingConvention: api schema model_path: target, + // biome-ignore lint/style/useNamingConvention: api schema hf_token: null, + // biome-ignore lint/style/useNamingConvention: api schema max_seq_length: isGguf ? 0 : 4096, + // biome-ignore lint/style/useNamingConvention: api schema load_in_4bit: true, + // biome-ignore lint/style/useNamingConvention: api schema is_lora: false, - gguf_variant: null, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: ggufVariant || null, + // biome-ignore lint/style/useNamingConvention: api schema trust_remote_code: false, + // biome-ignore lint/style/useNamingConvention: api schema chat_template_override: null, + // biome-ignore lint/style/useNamingConvention: api schema cache_type_kv: null, + // biome-ignore lint/style/useNamingConvention: api schema speculative_type: null, }); - toast.success(`Loaded ${target}`, { id: toastId, duration: 2000 }); + toast.success(`Loaded ${modelLabel}`, { id: toastId, duration: 2000 }); return null; } catch (error) { toast.dismiss(toastId); @@ -117,6 +273,147 @@ async function ensureLocalModelLoaded( } } +function getLocalModelLoadPlanForPayload( + payload: RecipePayload, +): LocalModelLoadPlan | null { + const localProviderNames = getLocalProviderNames(payload); + if (localProviderNames.size === 0) { + return null; + } + + const boundConfigs = findUsedLocalModelConfigs(payload, localProviderNames); + return getLocalModelLoadPlan(boundConfigs); +} + +async function getActiveLocalModelSelection(): Promise { + try { + const status = await getInferenceStatus(); + const target = status.active_model?.trim(); + if (!target) { + return null; + } + return { + target, + ggufVariant: status.gguf_variant?.trim() ?? "", + aliases: ["previous Chat model"], + }; + } catch { + return null; + } +} + +async function getRestorableActiveLocalModelSelection(): Promise { + try { + const status = await getInferenceStatus(); + const activeLabel = status.active_model?.trim() ?? null; + const target = ( + status.model_identifier ?? (status.is_gguf ? null : status.active_model) + )?.trim(); + if (!target) { + return { + selection: null, + unrestorableLabel: activeLabel, + }; + } + return { + selection: { + target, + ggufVariant: status.gguf_variant?.trim() ?? "", + aliases: ["previous Chat model"], + }, + unrestorableLabel: null, + }; + } catch { + return { selection: null, unrestorableLabel: null }; + } +} + +function isSameLocalModelSelection( + left: LocalModelSelection | null, + right: LocalModelSelection, +): boolean { + return Boolean( + left && + left.target.toLowerCase() === right.target.toLowerCase() && + left.ggufVariant === right.ggufVariant, + ); +} + +async function ensureLocalModelLoaded( + payload: RecipePayload, +): Promise { + const loadPlan = getLocalModelLoadPlanForPayload(payload); + if (!loadPlan) { + return null; + } + if (loadPlan.legacyAliases) { + const activeSelection = await getActiveLocalModelSelection(); + return activeSelection + ? null + : `Existing recipe uses legacy local model for ${loadPlan.legacyAliases.join(", ")}. Select a concrete local model or load one in Chat.`; + } + if (!loadPlan.selection) { + return loadPlan.error; + } + if (await isLocalModelAlreadyLoaded(loadPlan.selection)) { + return null; + } + return loadLocalModelSelection(loadPlan.selection); +} + +async function prepareLocalModelForRun(payload: RecipePayload): Promise<{ + error: string | null; + restorePrevious: (() => Promise) | null; +}> { + const loadPlan = getLocalModelLoadPlanForPayload(payload); + if (!loadPlan) { + return { error: null, restorePrevious: null }; + } + if (loadPlan.legacyAliases) { + const activeSelection = await getActiveLocalModelSelection(); + return activeSelection + ? { error: null, restorePrevious: null } + : { + error: `Existing recipe uses legacy local model for ${loadPlan.legacyAliases.join(", ")}. Select a concrete local model or load one in Chat.`, + restorePrevious: null, + }; + } + if (!loadPlan.selection) { + return { error: loadPlan.error, restorePrevious: null }; + } + if (await isLocalModelAlreadyLoaded(loadPlan.selection)) { + return { error: null, restorePrevious: null }; + } + + const previousSnapshot = await getRestorableActiveLocalModelSelection(); + const previousSelection = previousSnapshot.selection; + const error = await loadLocalModelSelection(loadPlan.selection); + if (error) { + return { error, restorePrevious: null }; + } + if (isSameLocalModelSelection(previousSelection, loadPlan.selection)) { + return { error: null, restorePrevious: null }; + } + return { + error: null, + restorePrevious: previousSelection + ? async () => { + const restoreError = await loadLocalModelSelection(previousSelection); + if (restoreError) { + toastError("Could not restore previous local model", restoreError); + } + } + : previousSnapshot.unrestorableLabel + ? () => { + toast.warning("Previous local model was not restored", { + description: `${previousSnapshot.unrestorableLabel} was selected from a native file path. Reopen it in Chat to continue with that model.`, + }); + return Promise.resolve(); + } + : null, + }; +} + type UseRecipeExecutionsParams = { recipeId: string; currentSignature: string; @@ -161,7 +458,11 @@ type UseRecipeExecutionsResult = { }; function formatValidationMessages(input: { - errors: Array<{ message: string; path?: string | null; code?: string | null }>; + errors: Array<{ + message: string; + path?: string | null; + code?: string | null; + }>; }): string[] { return input.errors.map((item) => { const path = item.path?.trim(); @@ -249,7 +550,8 @@ export function useRecipeExecutions({ (record: RecipeExecutionRecord): void => { const normalizedRecord = withExecutionDefaults(record); upsertExecution(normalizedRecord); - void saveRecipeExecution(normalizedRecord).catch((error) => { + saveRecipeExecution(normalizedRecord).catch((error) => { + // biome-ignore lint/suspicious/noConsole: background persistence failures should not interrupt the UI console.error("Save recipe execution failed:", error); }); }, @@ -287,7 +589,7 @@ export function useRecipeExecutions({ return; } - void trackRecipeExecution({ + trackRecipeExecution({ label: executionLabel(resumable.kind), kind: resumable.kind, rows: resumable.rows, @@ -299,11 +601,12 @@ export function useRecipeExecutions({ onPreviewSuccess, }); } catch (error) { + // biome-ignore lint/suspicious/noConsole: hydration failures are non-blocking diagnostics console.error("Load recipe executions failed:", error); } } - void hydrate(); + hydrate(); return () => { cancelled = true; @@ -344,9 +647,11 @@ export function useRecipeExecutions({ rows: number; settings: RecipeRunSettings; runName: string | null; + restorePrevious?: (() => Promise) | null; }): Promise => { - const { kind, payload, rows, settings, runName } = input; - const setLoading = kind === "preview" ? setPreviewLoading : setFullLoading; + const { kind, payload, rows, settings, runName, restorePrevious } = input; + const setLoading = + kind === "preview" ? setPreviewLoading : setFullLoading; const label = executionLabel(kind); setLoading(true); @@ -362,6 +667,8 @@ export function useRecipeExecutions({ onExecutionStart?.(); setRunDialogOpen(false); + let jobCreated = false; + let shouldRestorePrevious = false; try { const jobPayload = buildExecutionPayload({ payload, @@ -371,13 +678,14 @@ export function useRecipeExecutions({ runName, }); const createdJob = await createRecipeJob(jobPayload); + jobCreated = true; const executionWithJob = { ...baseExecution, jobId: createdJob.job_id, }; upsertAndPersist(executionWithJob); - return await trackRecipeExecution({ + const tracked = await trackRecipeExecution({ label, kind, rows, @@ -388,6 +696,8 @@ export function useRecipeExecutions({ onSetPreviewErrors: setRunErrors, onPreviewSuccess, }); + shouldRestorePrevious = tracked.terminal; + return tracked.success; } catch (error) { const message = toErrorMessage(error, `${label} request failed.`); upsertAndPersist({ @@ -398,8 +708,14 @@ export function useRecipeExecutions({ }); setRunErrors([message]); toastError(`${label} failed`, message); + if (!jobCreated) { + shouldRestorePrevious = true; + } return false; } finally { + if (shouldRestorePrevious && restorePrevious) { + await restorePrevious(); + } setLoading(false); } }, @@ -416,6 +732,48 @@ export function useRecipeExecutions({ ], ); + const prepareLocalModelForExecution = useCallback( + async ( + payload: RecipePayload, + ): Promise<(() => Promise) | null | false> => { + const { error, restorePrevious } = await prepareLocalModelForRun(payload); + if (!error) { + return restorePrevious; + } + setRunErrors([error]); + toastError("Local model failed to load", error); + return false; + }, + [setRunErrors], + ); + + const validateExecutionPayload = useCallback( + async ( + executionPayload: Parameters[0], + ): Promise => { + try { + const validation = await validateRecipe(executionPayload); + if (validation.valid) { + return true; + } + const errors = formatValidationMessages({ + errors: validation.errors, + }); + const fallback = validation.raw_detail ?? "Validation failed."; + const nextErrors = errors.length > 0 ? errors : [fallback]; + setRunErrors(nextErrors); + toastError("Validation failed", nextErrors[0]); + return false; + } catch (error) { + const message = toErrorMessage(error, "Validation failed."); + setRunErrors([message]); + toastError("Validation failed", message); + return false; + } + }, + [setRunErrors], + ); + const runWithValidation = useCallback( async ( kind: RecipeExecutionKind, @@ -435,20 +793,11 @@ export function useRecipeExecutions({ return false; } - // Flip to the Runs pane BEFORE we run ensureLocalModelLoaded + validate. - // Validation re-crawls the seed (multiple seconds for the github_repo - // reader) and the user otherwise stares at a "Running..." button with - // nothing else changing. runExecution() later no-ops this callback if - // the view has already been flipped, so we fire it once here. + // Flip to the Runs pane before validation starts. Validation can re-crawl + // the seed (multiple seconds for the github_repo reader), and runExecution() + // later no-ops this callback if the view has already been flipped. onExecutionStart?.(); - const localLoadError = await ensureLocalModelLoaded(payload); - if (localLoadError) { - setRunErrors([localLoadError]); - toastError("Local model failed to load", localLoadError); - return false; - } - const normalizedRows = sanitizeExecutionRows(rows, kind); const executionPayload = buildExecutionPayload({ payload, @@ -458,20 +807,17 @@ export function useRecipeExecutions({ runName, }); - try { - const validation = await validateRecipe(executionPayload); - if (!validation.valid) { - const errors = formatValidationMessages({ errors: validation.errors }); - const fallback = validation.raw_detail ?? "Validation failed."; - const nextErrors = errors.length > 0 ? errors : [fallback]; - setRunErrors(nextErrors); - toastError("Validation failed", nextErrors[0]); - return false; - } - } catch (error) { - const message = toErrorMessage(error, "Validation failed."); - setRunErrors([message]); - toastError("Validation failed", message); + if (!(await validateExecutionPayload(executionPayload))) { + return false; + } + + // Recipe and Chat share one singleton local inference backend. This + // direct load is a point-in-time handoff to job creation, not a lease: + // if Chat swaps models after this succeeds, the backend will reject or + // run against the active backend state. A future generation token should + // be validated across this load and the `/jobs` loaded-model gate. + const restorePrevious = await prepareLocalModelForExecution(payload); + if (restorePrevious === false) { return false; } @@ -481,26 +827,29 @@ export function useRecipeExecutions({ rows: normalizedRows, settings: runSettings, runName, + restorePrevious, }); }, [ onExecutionStart, + prepareLocalModelForExecution, readExecutablePayload, runExecution, runSettings, setRunErrors, + validateExecutionPayload, ], ); - const runPreview = useCallback(async (): Promise => { + const runPreview = useCallback((): Promise => { return runWithValidation("preview", previewRows, null); }, [previewRows, runWithValidation]); - const runFull = useCallback(async (): Promise => { + const runFull = useCallback((): Promise => { return runWithValidation("full", fullRows, fullRunName); }, [fullRows, fullRunName, runWithValidation]); - const runFromDialog = useCallback(async (): Promise => { + const runFromDialog = useCallback((): Promise => { setValidateResult(null); if (runDialogKind === "preview") { return runPreview(); @@ -512,9 +861,10 @@ export function useRecipeExecutions({ setRunErrors([]); const payload = readPayload(); if (!payload) { - const nextErrors = payloadResult.errors.length > 0 - ? payloadResult.errors - : [payloadErrorMessage]; + const nextErrors = + payloadResult.errors.length > 0 + ? payloadResult.errors + : [payloadErrorMessage]; setValidateResult({ valid: false, errors: nextErrors, @@ -525,24 +875,46 @@ export function useRecipeExecutions({ const rows = runDialogKind === "preview" ? previewRows : fullRows; const normalizedRows = sanitizeExecutionRows(rows, runDialogKind); - const executionPayload = buildExecutionPayload({ - payload, - kind: runDialogKind, - rows: normalizedRows, - settings: runSettings, - runName: runDialogKind === "full" ? normalizeRunName(fullRunName) : null, - }); setValidateLoading(true); try { + const executionPayload = buildExecutionPayload({ + payload, + kind: runDialogKind, + rows: normalizedRows, + settings: runSettings, + runName: + runDialogKind === "full" ? normalizeRunName(fullRunName) : null, + }); const validation = await validateRecipe(executionPayload); const errors = formatValidationMessages({ errors: validation.errors }); + if (!validation.valid) { + setValidateResult({ + valid: false, + errors, + rawDetail: validation.raw_detail ?? null, + }); + return false; + } + + const localLoadError = await ensureLocalModelLoaded(payload); + if (localLoadError) { + setRunErrors([localLoadError]); + setValidateResult({ + valid: false, + errors: [localLoadError], + rawDetail: null, + }); + toastError("Local model failed to load", localLoadError); + return false; + } + setValidateResult({ - valid: validation.valid, + valid: true, errors, rawDetail: validation.raw_detail ?? null, }); - return validation.valid; + return true; } catch (error) { const message = toErrorMessage(error, "Validation failed."); setValidateResult({ @@ -612,7 +984,12 @@ export function useRecipeExecutions({ const loadExecutionDatasetPage = useCallback( async (id: string, page: number): Promise => { const execution = executions.find((entry) => entry.id === id); - if (!execution || execution.kind !== "full" || !execution.jobId || page < 1) { + if ( + !execution || + execution.kind !== "full" || + !execution.jobId || + page < 1 + ) { return; } @@ -625,7 +1002,9 @@ export function useRecipeExecutions({ }); const dataset = normalizeDatasetRows(response.dataset); const total = - typeof response.total === "number" ? response.total : execution.datasetTotal; + typeof response.total === "number" + ? response.total + : execution.datasetTotal; upsertAndPersist({ ...execution, dataset, diff --git a/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts b/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts index 56634b84c8..dd5e12233c 100644 --- a/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts +++ b/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts @@ -97,7 +97,9 @@ export function applyRenameToConfig( next = { ...base, // biome-ignore lint/style/useNamingConvention: api schema - target_columns: targets.map((target) => (target === from ? to : target)), + target_columns: targets.map((target) => + target === from ? to : target, + ), }; } } @@ -137,14 +139,12 @@ export function applyRemovalToConfig( } if (config.kind === "model_config" && config.provider === ref) { const base = next as ModelConfig; - // Clear the synthetic "local" placeholder when the provider that was - // a local provider is removed; otherwise the stale placeholder would - // pass validation against a future external provider and then fail - // at runtime against a real API ("model not found"). next = { ...base, provider: "", - model: base.model === "local" ? "" : base.model, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, }; } if (config.kind === "llm" && config.model_alias === ref) { @@ -156,7 +156,9 @@ export function applyRemovalToConfig( next = { ...base, tool_alias: "" }; } if (config.kind === "validator") { - const targets = (config.target_columns ?? []).filter((target) => target !== ref); + const targets = (config.target_columns ?? []).filter( + (target) => target !== ref, + ); if (targets.length !== (config.target_columns ?? []).length) { const base = next as typeof config; next = { @@ -206,5 +208,7 @@ export function applyRemovalToConfigs( if (!ref) { return configs; } - return applyConfigTransform(configs, (config) => applyRemovalToConfig(config, ref)); + return applyConfigTransform(configs, (config) => + applyRemovalToConfig(config, ref), + ); } diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 036cc94082..52ce1329bf 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -12,23 +12,23 @@ import { applyNodeChanges, } from "@xyflow/react"; import { create } from "zustand"; -import type { - RecipeNode, - RecipeProcessorConfig, - LayoutDirection, - LlmType, - NodeConfig, - SeedSourceType, - SamplerType, -} from "../types"; import { - getBlockDefinition, type BlockKind, type BlockType, type SeedBlockType, + getBlockDefinition, } from "../blocks/registry"; -import { deriveDisplayGraph } from "../utils/graph/derive-display-graph"; +import type { + LayoutDirection, + LlmType, + NodeConfig, + RecipeNode, + RecipeProcessorConfig, + SamplerType, + SeedSourceType, +} from "../types"; import { applyRecipeConnection, isValidRecipeConnection } from "../utils/graph"; +import { deriveDisplayGraph } from "../utils/graph/derive-display-graph"; import { HANDLE_IDS, normalizeRecipeHandleId, @@ -42,8 +42,8 @@ import { } from "./helpers/model-infra-layout"; import { applyEdgeRemovals, applyNodeRemovals } from "./helpers/removals"; import { - applyRenameToConfigs, applyLayoutDirectionToNodes, + applyRenameToConfigs, buildNodeUpdate, syncEdgesForConfigPatch, syncSubcategoryConfigsForCategoryUpdate, @@ -97,7 +97,11 @@ type RecipeStudioState = { position?: XYPosition, openDialog?: boolean, ) => void; - addLlmNode: (type: LlmType, position?: XYPosition, openDialog?: boolean) => void; + addLlmNode: ( + type: LlmType, + position?: XYPosition, + openDialog?: boolean, + ) => void; addModelProviderNode: (position?: XYPosition, openDialog?: boolean) => void; addModelConfigNode: (position?: XYPosition, openDialog?: boolean) => void; addToolProfileNode: (position?: XYPosition, openDialog?: boolean) => void; @@ -250,7 +254,10 @@ function connectSemantic( }; } -function isModelSemanticEdge(edge: Edge, configs: Record): boolean { +function isModelSemanticEdge( + edge: Edge, + configs: Record, +): boolean { const source = configs[edge.source]; const target = configs[edge.target]; return Boolean( @@ -315,12 +322,16 @@ export const useRecipeStudioStore = create((set, get) => ({ auxNodePositions: {}, llmAuxVisibility: state.llmAuxVisibility, }); - const { nodes } = getLayoutedElements(displayGraph.nodes, displayGraph.edges, { - direction: state.layoutDirection, - nodesep: isTopBottom ? 120 : 80, - ranksep: isTopBottom ? 140 : 80, - configs: state.configs, - }); + const { nodes } = getLayoutedElements( + displayGraph.nodes, + displayGraph.edges, + { + direction: state.layoutDirection, + nodesep: isTopBottom ? 120 : 80, + ranksep: isTopBottom ? 140 : 80, + configs: state.configs, + }, + ); const layoutedPositions = new Map( nodes.map((node) => [node.id, node.position] as const), ); @@ -381,13 +392,7 @@ export const useRecipeStudioStore = create((set, get) => ({ (config) => config.kind === "seed", ); if (!existing) { - return buildAddedNodeState( - state, - "seed", - type, - position, - openDialog, - ); + return buildAddedNodeState(state, "seed", type, position, openDialog); } let nextSourceType: SeedSourceType = "hf"; if (type === "seed_local") { @@ -430,7 +435,10 @@ export const useRecipeStudioStore = create((set, get) => ({ [existing.id]: nextConfig, }, nodes: updateNodeData( - state.nodes.map((node) => ({ ...node, selected: node.id === existing.id })), + state.nodes.map((node) => ({ + ...node, + selected: node.id === existing.id, + })), existing.id, nextConfig, state.layoutDirection, @@ -444,7 +452,13 @@ export const useRecipeStudioStore = create((set, get) => ({ if (state.executionLocked) { return state; } - const added = buildAddedNodeState(state, "llm", type, position, openDialog); + const added = buildAddedNodeState( + state, + "llm", + type, + position, + openDialog, + ); const context = getAddedNodeContext(added); if (!context) { return added; @@ -495,9 +509,7 @@ export const useRecipeStudioStore = create((set, get) => ({ let { nodes, configs } = context; let edges = state.edges; const unboundModelConfigs = Object.values(configs).filter( - (config) => - config.kind === "model_config" && - !config.provider.trim(), + (config) => config.kind === "model_config" && !config.provider.trim(), ); if (!position && unboundModelConfigs.length > 0) { nodes = placeNodeNear( @@ -605,7 +617,7 @@ export const useRecipeStudioStore = create((set, get) => ({ let { nodes, configs } = context; let edges = state.edges; const unboundLlms = Object.values(configs).filter( - (config) => config.kind === "llm" && !(config.tool_alias?.trim()), + (config) => config.kind === "llm" && !config.tool_alias?.trim(), ); if (!position && unboundLlms.length > 0) { nodes = placeNodeNear( @@ -757,17 +769,15 @@ export const useRecipeStudioStore = create((set, get) => ({ if (cfg.kind !== "model_config" || cfg.provider !== providerName) { continue; } - if (nextIsLocal && !cfg.model.trim()) { - // external -> local: auto fill the placeholder model id so the - // config does not fail "model is required" validation. - configs = { ...configs, [cfgId]: { ...cfg, model: "local" } }; - continue; - } - if (!nextIsLocal && cfg.model === "local") { - // local -> external: clear the placeholder so the user picks a - // real model id for the new external endpoint. - configs = { ...configs, [cfgId]: { ...cfg, model: "" } }; - } + configs = { + ...configs, + [cfgId]: { + ...cfg, + model: "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + }, + }; } } } diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index 9c720a06d3..b8ed13f70b 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -264,6 +264,8 @@ export type ModelConfig = { kind: "model_config"; name: string; model: string; + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant?: string; provider: string; // biome-ignore lint/style/useNamingConvention: api schema inference_temperature?: string; diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts b/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts index 4fa70b3e3a..059ddf186f 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts @@ -11,7 +11,6 @@ import { isSemanticTargetHandle, normalizeRecipeHandleId, } from "../handles"; -import { isSemanticRelation } from "./relations"; import { isCategoryConfig, isExpressionConfig, @@ -21,6 +20,7 @@ import { VALIDATOR_OXC_CODE_LANGS, VALIDATOR_SQL_CODE_LANGS, } from "../validators/code-lang"; +import { isSemanticRelation } from "./relations"; function buildTemplateWithRef(template: string, ref: string): string { if (template.includes(ref)) { @@ -157,7 +157,10 @@ function isCompetingIncomingEdge( return source.kind === "sampler" && source.sampler_type === "datetime"; } -function isModelSemanticRelation(source: NodeConfig, target: NodeConfig): boolean { +function isModelSemanticRelation( + source: NodeConfig, + target: NodeConfig, +): boolean { return ( (source.kind === "model_provider" && target.kind === "model_config") || (source.kind === "model_config" && target.kind === "llm") || @@ -181,7 +184,9 @@ function canApplyCodeLangToValidator( if (normalized === "python") { return true; } - return VALIDATOR_SQL_CODE_LANGS.includes(normalized as typeof validator.code_lang); + return VALIDATOR_SQL_CODE_LANGS.includes( + normalized as typeof validator.code_lang, + ); } function countHandleUsage( @@ -333,12 +338,8 @@ export function applyRecipeConnection( if (!isValidRecipeConnection(connection, configs)) { return { edges }; } - const initialSource = connection.source - ? configs[connection.source] - : null; - const initialTarget = connection.target - ? configs[connection.target] - : null; + const initialSource = connection.source ? configs[connection.source] : null; + const initialTarget = connection.target ? configs[connection.target] : null; if (!(initialSource && initialTarget)) { return { edges }; } @@ -386,17 +387,36 @@ export function applyRecipeConnection( nextBaseEdges, ); if (source.kind === "model_provider" && target.kind === "model_config") { - // Keep the model_config.model field in sync with provider mode when the - // link is changed via graph drag (the model-config dialog path has its - // own applyProviderChange helper that does the same thing). + // Keep model_config.provider in sync when a graph drag changes the link. + // Local providers now require an explicit selected load id; do not synthesize + // the legacy "local" placeholder. External relinks clear local-only GGUF + // metadata, while legacy placeholders are normalized back to empty. const isSourceLocal = source.is_local === true; - let nextModel = target.model; - if (isSourceLocal && !nextModel.trim()) { - nextModel = "local"; - } else if (!isSourceLocal && nextModel === "local") { - nextModel = ""; - } - const next = { ...target, provider: source.name, model: nextModel }; + const isLegacyLocalPlaceholder = + target.model.trim().toLowerCase() === "local"; + const previousProviderName = target.provider.trim(); + const previousProvider = Object.values(configs).find( + (config) => + config.kind === "model_provider" && + config.name === previousProviderName, + ); + const wasLinkedToLocal = + previousProvider?.kind === "model_provider" && + previousProvider.is_local === true; + const shouldClearModel = + isLegacyLocalPlaceholder || + (isSourceLocal ? !wasLinkedToLocal : wasLinkedToLocal); + const next = { + ...target, + provider: source.name, + ...(shouldClearModel ? { model: "" } : {}), + ...(shouldClearModel || !isSourceLocal + ? { + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: undefined, + } + : {}), + }; return { edges: nextEdges, configs: { ...configs, [target.id]: next } }; } if (source.kind === "model_config" && target.kind === "llm") { @@ -435,10 +455,9 @@ export function applyRecipeConnection( // biome-ignore lint/style/useNamingConvention: api schema target_columns: [source.name], // biome-ignore lint/style/useNamingConvention: api schema - code_lang: - ( - canUseCodeLangForTarget ? nextCodeLang : target.code_lang - ) as typeof target.code_lang, + code_lang: (canUseCodeLangForTarget + ? nextCodeLang + : target.code_lang) as typeof target.code_lang, }; return { edges: nextEdges, configs: { ...configs, [target.id]: next } }; } diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts index 15ecf39a7b..6b3df846a1 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts @@ -1,15 +1,8 @@ // 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 type { - ModelConfig, - ModelProviderConfig, -} from "../../../types"; -import { - isRecord, - readNumberString, - readString, -} from "../helpers"; +import type { ModelConfig, ModelProviderConfig } from "../../../types"; +import { isRecord, readNumberString, readString } from "../helpers"; export function parseModelProvider( provider: Record, @@ -53,6 +46,8 @@ export function parseModelConfig( kind: "model_config", name, model: readString(model.model) ?? "", + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: readString(model.gguf_variant) ?? undefined, provider: readString(model.provider) ?? "", // biome-ignore lint/style/useNamingConvention: api schema inference_temperature: readNumberString(inference.temperature), diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts index 14e0faa5cc..1575919705 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts @@ -54,55 +54,62 @@ export function buildModelProvider( }; } -export function buildModelConfig( +function assignFiniteNumber( + target: Record, + key: string, + rawValue: string | undefined, + transform: (value: number) => number = (value) => value, +): void { + const trimmed = rawValue?.trim(); + if (!trimmed) { + return; + } + + const parsed = Number(trimmed); + if (Number.isFinite(parsed)) { + target[key] = transform(parsed); + } +} + +function buildInferenceParameters( config: ModelConfig, errors: string[], ): Record { const inference: Record = {}; - const temp = config.inference_temperature?.trim(); - const topP = config.inference_top_p?.trim(); - const maxTokens = config.inference_max_tokens?.trim(); - const timeout = config.inference_timeout?.trim(); + assignFiniteNumber(inference, "temperature", config.inference_temperature); + assignFiniteNumber(inference, "top_p", config.inference_top_p); + assignFiniteNumber(inference, "max_tokens", config.inference_max_tokens); + assignFiniteNumber( + inference, + "timeout", + config.inference_timeout, + Math.trunc, + ); + const extraBody = parseJsonObject( config.inference_extra_body, `Model ${config.name} inference extra_body`, errors, ); - - if (temp) { - const parsed = Number(temp); - if (Number.isFinite(parsed)) { - inference.temperature = parsed; - } - } - if (topP) { - const parsed = Number(topP); - if (Number.isFinite(parsed)) { - // biome-ignore lint/style/useNamingConvention: api schema - inference.top_p = parsed; - } - } - if (maxTokens) { - const parsed = Number(maxTokens); - if (Number.isFinite(parsed)) { - // biome-ignore lint/style/useNamingConvention: api schema - inference.max_tokens = parsed; - } - } - if (timeout) { - const parsed = Number(timeout); - if (Number.isFinite(parsed)) { - inference.timeout = Math.trunc(parsed); - } - } if (extraBody) { - // biome-ignore lint/style/useNamingConvention: api schema inference.extra_body = extraBody; } + return inference; +} + +export function buildModelConfig( + config: ModelConfig, + errors: string[], +): Record { + const inference = buildInferenceParameters(config, errors); + const ggufVariant = config.gguf_variant?.trim(); + return { alias: config.name, model: config.model, + // biome-ignore lint/style/useNamingConvention: api schema + gguf_variant: ggufVariant || undefined, provider: config.provider || undefined, // biome-ignore lint/style/useNamingConvention: api schema inference_parameters: diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts b/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts index 7e72e8d919..a3b3763291 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts @@ -54,7 +54,9 @@ export function validateTimedeltaConfigs( } const reference = config.reference_column_name?.trim() ?? ""; if (!reference) { - errors.push(`Timedelta ${config.name}: reference datetime column required.`); + errors.push( + `Timedelta ${config.name}: reference datetime column required.`, + ); continue; } const parent = nameToConfig.get(reference); @@ -63,7 +65,9 @@ export function validateTimedeltaConfigs( parent.kind !== "sampler" || parent.sampler_type !== "datetime" ) { - errors.push(`Timedelta ${config.name}: reference '${reference}' must be datetime.`); + errors.push( + `Timedelta ${config.name}: reference '${reference}' must be datetime.`, + ); } } } @@ -91,9 +95,18 @@ export function validateModelConfigProviders( const provider = config.provider.trim(); const alias = config.name; const isLocal = localProviderNames.has(provider); - // Local providers do not require a real model id - the loaded Chat - // model is used regardless of what gets sent in the payload. - if (!isLocal && modelAliases.has(alias) && !config.model.trim()) { + const isUsed = modelAliases.has(alias); + const model = config.model.trim(); + const isLegacyLocalPlaceholder = model.toLowerCase() === "local"; + + if (!isLocal && isUsed && isLegacyLocalPlaceholder) { + errors.push(`Model config ${alias}: model is required.`); + continue; + } + if (isLocal && isUsed && !model) { + errors.push(`Model config ${alias}: choose a local model.`); + } + if (!isLocal && isUsed && !model) { errors.push(`Model config ${alias}: model is required.`); } if (provider && !modelProviderNames.has(provider)) { @@ -121,7 +134,9 @@ export function validateUsedProviders( errors.push(`Model provider ${provider.name}: endpoint is required.`); } if (!provider.provider_type.trim()) { - errors.push(`Model provider ${provider.name}: provider_type is required.`); + errors.push( + `Model provider ${provider.name}: provider_type is required.`, + ); } } } @@ -145,7 +160,9 @@ export function validateValidatorConfigs( continue; } if (targetConfig.kind !== "llm" || targetConfig.llm_type !== "code") { - errors.push(`Validator ${config.name}: target '${target}' must be LLM Code.`); + errors.push( + `Validator ${config.name}: target '${target}' must be LLM Code.`, + ); continue; } if ( diff --git a/studio/frontend/src/features/settings/components/api-key-row.tsx b/studio/frontend/src/features/settings/components/api-key-row.tsx index ced6de17d1..9a6e1ee207 100644 --- a/studio/frontend/src/features/settings/components/api-key-row.tsx +++ b/studio/frontend/src/features/settings/components/api-key-row.tsx @@ -14,30 +14,39 @@ import { MoreHorizontalIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; +import { useT } from "@/i18n"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import type { ApiKey } from "../api/api-keys"; -function relative(iso: string | null): string { - if (!iso) return "never"; +type SettingsT = ReturnType; + +function relative(iso: string | null, t: SettingsT): string { + if (!iso) return t("settings.apiKeys.relativeNever"); const diff = Date.now() - new Date(iso).getTime(); const days = Math.floor(diff / 86400000); if (days < 1) { const hours = Math.floor(diff / 3600000); - if (hours < 1) return "just now"; - return `${hours}h ago`; + if (hours < 1) return t("settings.apiKeys.relativeJustNow"); + return t("settings.apiKeys.relativeHoursAgo", { count: hours }); } - if (days < 30) return `${days}d ago`; - if (days < 365) return `${Math.floor(days / 30)}mo ago`; - return `${Math.floor(days / 365)}y ago`; + if (days < 30) return t("settings.apiKeys.relativeDaysAgo", { count: days }); + if (days < 365) { + return t("settings.apiKeys.relativeMonthsAgo", { + count: Math.floor(days / 30), + }); + } + return t("settings.apiKeys.relativeYearsAgo", { + count: Math.floor(days / 365), + }); } -function expiresText(iso: string | null): string { - if (!iso) return "never"; +function expiresText(iso: string | null, t: SettingsT): string { + if (!iso) return t("settings.apiKeys.relativeNever"); const diff = new Date(iso).getTime() - Date.now(); - if (diff < 0) return "expired"; + if (diff < 0) return t("settings.apiKeys.expired"); const days = Math.floor(diff / 86400000); - if (days < 1) return "today"; - return `in ${days}d`; + if (days < 1) return t("settings.apiKeys.today"); + return t("settings.apiKeys.inDays", { count: days }); } export function ApiKeyRow({ @@ -47,6 +56,7 @@ export function ApiKeyRow({ apiKey: ApiKey; onRevoke: (key: ApiKey) => void; }) { + const t = useT(); const prefix = `sk-unsloth-${apiKey.key_prefix}…`; return (
@@ -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), + })} +
@@ -77,7 +99,7 @@ export function ApiKeyRow({ variant="ghost" size="sm" className="size-7 p-0 opacity-0 transition-opacity group-hover:opacity-100 data-[state=open]:opacity-100 max-sm:!opacity-100 max-sm:size-9" - aria-label={`Actions for ${apiKey.name}`} + aria-label={t("settings.apiKeys.actionsFor", { name: apiKey.name })} > @@ -85,14 +107,14 @@ export function ApiKeyRow({ { await copyToClipboard(prefix); }}> - Copy prefix + {t("settings.apiKeys.copyPrefix")} onRevoke(apiKey)} className="text-destructive focus:text-destructive" > - Revoke token + {t("settings.apiKeys.revokeToken")} diff --git a/studio/frontend/src/features/settings/components/create-key-form.tsx b/studio/frontend/src/features/settings/components/create-key-form.tsx index a0f2d7f82d..93802bfad8 100644 --- a/studio/frontend/src/features/settings/components/create-key-form.tsx +++ b/studio/frontend/src/features/settings/components/create-key-form.tsx @@ -3,6 +3,7 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; +import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { useState } from "react"; import { createApiKey } from "../api/api-keys"; @@ -21,6 +22,7 @@ export function CreateKeyForm({ onCreated: (rawKey: string) => void; onError: (message: string) => void; }) { + const t = useT(); const [name, setName] = useState(""); const [expiry, setExpiry] = useState(null); const [loading, setLoading] = useState(false); @@ -33,8 +35,11 @@ export function CreateKeyForm({ const result = await createApiKey(name.trim(), expiry); onCreated(result.key); setName(""); - } catch (err) { - onError(err instanceof Error ? err.message : "Couldn't create access token."); + } catch { + // API helpers in ../api/api-keys.ts throw generic English Error + // messages; always use the translated message so zh-CN users do not + // see English text bleed through from internal exceptions. + onError(t("settings.apiKeys.createError")); } finally { setLoading(false); } @@ -49,9 +54,9 @@ export function CreateKeyForm({ setName(e.target.value)} - placeholder="Token name (e.g. production)" + placeholder={t("settings.apiKeys.tokenNamePlaceholder")} className="h-8 min-w-[180px] flex-1 text-sm" - aria-label="New access token name" + aria-label={t("settings.apiKeys.newAccessTokenName")} />
{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")}

diff --git a/studio/frontend/src/features/settings/components/language-select.tsx b/studio/frontend/src/features/settings/components/language-select.tsx new file mode 100644 index 0000000000..9d30e06147 --- /dev/null +++ b/studio/frontend/src/features/settings/components/language-select.tsx @@ -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 + +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { + LOCALES, + isSupportedLocale, + setLocale, + useT, + useLocale, +} from "@/i18n"; + +export function LanguageSelect() { + const t = useT(); + const locale = useLocale(); + + return ( + + ); +} diff --git a/studio/frontend/src/features/settings/components/theme-segmented.tsx b/studio/frontend/src/features/settings/components/theme-segmented.tsx index 1061995346..36e7062d8f 100644 --- a/studio/frontend/src/features/settings/components/theme-segmented.tsx +++ b/studio/frontend/src/features/settings/components/theme-segmented.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { cn } from "@/lib/utils"; +import { useT, type TranslationKey } from "@/i18n"; import { LaptopIcon, Moon02Icon, @@ -11,13 +12,18 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { motion, useReducedMotion } from "motion/react"; import { useTheme, type Theme } from "../stores/theme-store"; -const OPTIONS: { value: Theme; label: string; icon: typeof Sun02Icon }[] = [ - { value: "light", label: "Light", icon: Sun02Icon }, - { value: "dark", label: "Dark", icon: Moon02Icon }, - { value: "system", label: "System", icon: LaptopIcon }, +const OPTIONS: { + value: Theme; + labelKey: TranslationKey; + icon: typeof Sun02Icon; +}[] = [ + { value: "light", labelKey: "settings.appearance.theme.light", icon: Sun02Icon }, + { value: "dark", labelKey: "settings.appearance.theme.dark", icon: Moon02Icon }, + { value: "system", labelKey: "settings.appearance.theme.system", icon: LaptopIcon }, ]; export function ThemeSegmented() { + const t = useT(); const { theme, setTheme } = useTheme(); const reduced = useReducedMotion(); return ( @@ -49,7 +55,7 @@ export function ThemeSegmented() { /> )} - {opt.label} + {t(opt.labelKey)} ); })} diff --git a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx index e4cdccd2d7..90f66483b7 100644 --- a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx +++ b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -29,10 +30,13 @@ export type UpdateInstallSource = | "unknown"; type UpdateInstallSourceState = UpdateInstallSource | "loading"; -function getStudioUpdateInstructionLine(shell: UpdateShell): string { +function getStudioUpdateInstructionLine( + shell: UpdateShell, + t: ReturnType, +): string { return shell === "windows" - ? "Open PowerShell and run:" - : "Open Terminal and run:"; + ? t("settings.about.update.openPowerShell") + : t("settings.about.update.openTerminal"); } function isLocalInstallSource( @@ -59,6 +63,7 @@ function CopyableCommand({ command: string; copyLabel: string; }): ReactElement { + const t = useT(); const [copied, setCopied] = useState(false); const timerRef = useRef | null>(null); @@ -89,14 +94,26 @@ function CopyableCommand({ value={command} className="min-w-0 flex-1 bg-transparent px-2 py-1.5 font-mono text-[11px] text-foreground outline-none" title={command} - aria-label={`${copyLabel} text`} + aria-label={t("settings.about.update.commandText", { + label: copyLabel, + })} /> ); })} @@ -125,20 +131,20 @@ export function UsageExamples() { type="button" onClick={handleCopy} className="flex items-center gap-1 rounded px-1.5 py-1 text-[11px] text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring" - aria-label="Copy snippet" + aria-label={t("settings.apiKeys.copySnippet")} > - {copied ? "Copied" : "Copy"} + {copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")} 
           {snippets[lang]}
- Setup docs: + {t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( s.open); const activeTab = useSettingsDialogStore((s) => s.activeTab); const setActiveTab = useSettingsDialogStore((s) => s.setActiveTab); @@ -128,9 +144,9 @@ export function SettingsDialog() { "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", )} > - Settings + {t("settings.dialog.title")} - Manage your Unsloth Studio preferences. + {t("settings.dialog.description")}