ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-05-26 07:26:37 -07:00
commit d3ac7447eb
71 changed files with 9671 additions and 1047 deletions
Download patch file Download diff file Expand all files Collapse all files

2
.pre-commit-config.yaml
View file

@ -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:

37
install.ps1
View file

@ -1300,7 +1300,7 @@ 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.7" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -1314,7 +1314,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.7" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -1352,7 +1352,7 @@ 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.7" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
@ -1364,7 +1364,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.7" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -1392,7 +1392,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.7" --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)
@ -1626,6 +1626,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)

56
install.sh
View file

@ -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"
@ -1865,7 +1873,7 @@ 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
"unsloth>=2026.5.7" 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.
@ -1878,7 +1886,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.7" unsloth-zoo
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2046,7 +2054,7 @@ 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
"unsloth>=2026.5.7" 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
@ -2064,7 +2072,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.7" 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..."
@ -2096,7 +2104,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.7" --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..."
@ -2108,12 +2116,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.
@ -2263,6 +2265,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"

41
studio/backend/core/data_recipe/service.py
View file

@ -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(

2
studio/backend/core/export/export.py
View file

@ -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,

960
studio/backend/core/inference/external_provider.py
View file

File diff suppressed because it is too large Load diff

4
studio/backend/core/inference/llama_cpp.py
View file

@ -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))"

245
studio/backend/core/inference/pricing.py
View file

@ -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,
},

109
studio/backend/main.py
View file

@ -49,6 +49,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}$")
@ -715,10 +717,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
@ -743,6 +743,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():
@ -753,11 +833,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(
@ -767,11 +854,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"}
@ -785,6 +872,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

52
studio/backend/models/inference.py
View file

@ -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),
@ -786,6 +826,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":

5
studio/backend/requirements/single-env/overrides-darwin-arm64.txt Normal file
View file

@ -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

121
studio/backend/routes/data_recipe/jobs.py
View file

@ -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

70
studio/backend/routes/inference.py
View file

@ -606,11 +606,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).
@ -619,7 +624,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)
@ -1373,6 +1379,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,
@ -1435,6 +1442,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,
@ -1709,6 +1717,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
@ -1717,6 +1731,7 @@ def _build_external_messages(
"""
document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS
anthropic = provider_type == "anthropic"
openai = provider_type == "openai"
result = []
for msg in messages:
if isinstance(msg.content, str):
@ -1737,6 +1752,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'
@ -1758,6 +1797,8 @@ def _build_external_messages(
# provider would 400 on the unknown part, so
# gate by provider_type.
parts.append({"type": "compaction", "content": part.content})
if msg.role == "assistant" and not parts:
continue
result.append({"role": msg.role, "content": parts})
else:
# Non-vision provider: strip images / documents, keep
@ -1769,8 +1810,28 @@ 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.
@ -1876,6 +1937,7 @@ 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,
fast_mode = payload.fast_mode,
stream = payload.stream,
)
try:

133
studio/backend/run.py
View file

@ -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"
@ -512,10 +513,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 +669,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 <path/to/studio/frontend/dist>\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 +837,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")

353
studio/backend/tests/test_anthropic_citations.py Normal file
View file

@ -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

690
studio/backend/tests/test_anthropic_citations_edge.py Normal file
View file

@ -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("…")

164
studio/backend/tests/test_anthropic_fast_mode_and_refusal.py Normal file
View file

@ -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

442
studio/backend/tests/test_anthropic_fast_mode_edge.py Normal file
View file

@ -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

51
studio/backend/tests/test_anthropic_web_fetch.py
View file

@ -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 ─────────────────────────────────────────────────
@ -365,7 +351,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 +368,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

248
studio/backend/tests/test_frontend_resolution.py Normal file
View file

@ -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("<!doctype html>", 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("<!doctype html>", 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("<!doctype html>", 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("<!doctype html>", 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("<!doctype html>", 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("<!doctype html>", 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 <path/to/studio/frontend/dist>\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

144
studio/backend/tests/test_index_bootstrap_origin.py Normal file
View file

@ -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

196
studio/backend/tests/test_index_bootstrap_origin_extra.py Normal file
View file

@ -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,<script>alert(1)</script>"
)
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

6
studio/backend/tests/test_multimodal_document.py
View file

@ -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},
}

251
studio/backend/tests/test_openai_citation_markers.py Normal file
View file

@ -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

413
studio/backend/tests/test_openai_citation_markers_edge.py Normal file
View file

@ -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<sid>[\\ue202<sid>...][\\ue202<loc>]\\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<sid>`` 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<sid>`` 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

1
studio/backend/tests/test_openai_image_generation.py
View file

@ -210,6 +210,7 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch):
assert starts[0]["arguments"] == {
"kind": "image",
"prompt": "A photorealistic cat sitting",
"openai_image_generation_call_id": "img_abc",
}
assert ends[0]["image_b64"] == "AAAA"
assert ends[0]["image_mime"] == "image/png"

372
studio/backend/tests/test_openai_tool_result_fallbacks.py Normal file
View file

@ -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: <query>") 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: <query>` 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"]

281
studio/backend/tests/test_pricing.py
View file

@ -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

475
studio/backend/tests/test_pricing_edge.py Normal file
View file

@ -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

15
studio/backend/utils/hardware/hardware.py
View file

@ -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,

4
studio/frontend/package.json
View file

@ -12,8 +12,8 @@
"lint": "eslint .",
"preview": "vite preview",
"typecheck": "tsc -b --pretty false",
"biome:check": "biome check .",
"biome:fix": "biome check . --write"
"biome:check": "biome check",
"biome:fix": "biome check --write"
},
"dependencies": {
"@assistant-ui/core": "0.1.17",

77
studio/frontend/src/components/assistant-ui/generated-image-overlay-context.tsx Normal file
View file

@ -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<GeneratedImageOverlayContextValue | null>(null);
export function GeneratedImageOverlayProvider({
children,
threadId = null,
}: {
children: ReactNode;
threadId?: string | null;
}) {
const [overlay, setOverlay] = useState<GeneratedImageOverlayState | null>(
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 (
<GeneratedImageOverlayContext.Provider value={value}>
{children}
</GeneratedImageOverlayContext.Provider>
);
}
export function useGeneratedImageOverlay(): GeneratedImageOverlayContextValue {
const context = useContext(GeneratedImageOverlayContext);
if (!context) {
throw new Error(
"useGeneratedImageOverlay must be used within GeneratedImageOverlayProvider.",
);
}
return context;
}

510
studio/frontend/src/components/assistant-ui/image.tsx Normal file
View file

@ -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<ImageMessagePart, "image" | "filename">,
): 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<ImageMessagePart, "image">,
): Promise<void> => {
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<typeof imageVariants>;
function ImageRoot({
className,
variant,
size,
children,
...props
}: ImageRootProps) {
return (
<div
data-slot="image-root"
data-variant={variant}
data-size={size}
className={cn(imageVariants({ variant, size, className }))}
{...props}
>
{children}
</div>
);
}
type ImagePreviewProps = Omit<ComponentProps<"img">, "children"> & {
containerClassName?: string;
};
function ImagePreview({
className,
containerClassName,
onLoad,
onError,
alt = "Image content",
src,
...props
}: ImagePreviewProps) {
const imgRef = useRef<HTMLImageElement>(null);
const [loadedSrc, setLoadedSrc] = useState<string | undefined>(undefined);
const [errorSrc, setErrorSrc] = useState<string | undefined>(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 (
<div
data-slot="image-preview"
className={cn("relative min-h-32", containerClassName)}
>
{!(loaded || error) && (
<div
data-slot="image-preview-loading"
className="absolute inset-0 flex items-center justify-center bg-muted/50"
>
<ImageIcon className="size-8 animate-pulse text-muted-foreground" />
</div>
)}
{error ? (
<div
data-slot="image-preview-error"
className="flex min-h-32 items-center justify-center bg-muted/50 p-4"
>
<ImageOffIcon className="size-8 text-muted-foreground" />
</div>
) : (
<img
{...props}
ref={imgRef}
src={src}
alt={alt}
className={cn(
"block h-auto w-full object-contain",
!loaded && "invisible",
className,
)}
onLoad={(e) => {
if (typeof src === "string") {
setLoadedSrc(src);
}
onLoad?.(e);
}}
onError={(e) => {
if (typeof src === "string") {
setErrorSrc(src);
}
onError?.(e);
}}
/>
)}
</div>
);
}
function ImageFilename({
className,
children,
...props
}: ComponentProps<"span">) {
if (!children) {
return null;
}
return (
<span
data-slot="image-filename"
className={cn(
"block truncate px-2 py-1.5 text-muted-foreground text-xs",
className,
)}
{...props}
>
{children}
</span>
);
}
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 (
<>
<button
type="button"
onClick={handleOpen}
className="aui-image-zoom-trigger w-full cursor-zoom-in border-0 bg-transparent p-0 text-left"
aria-label="Click to zoom image"
>
{children}
</button>
{isOpen &&
typeof document !== "undefined" &&
createPortal(
<button
type="button"
data-slot="image-zoom-overlay"
className="aui-image-zoom-overlay fade-in fixed inset-0 z-50 flex animate-in items-center justify-center border-0 bg-black/80 p-0 duration-200"
onClick={handleClose}
aria-label="Close zoomed image"
>
<img
data-slot="image-zoom-content"
src={src}
alt={alt}
className="aui-image-zoom-content fade-in zoom-in-95 max-h-[90vh] max-w-[90vw] animate-in cursor-zoom-out object-contain duration-200"
/>
</button>,
document.body,
)}
</>
);
}
function ImageGenerating({ className }: { className?: string }) {
return (
<div
data-slot="image-generating"
className={cn(
"flex min-h-32 items-center justify-center bg-muted/50 p-4",
className,
)}
>
<Loader2Icon className="size-8 animate-spin text-muted-foreground" />
<span className="sr-only">Generating image…</span>
</div>
);
}
function ImageContentFilterError({
className,
reason,
}: {
className?: string;
reason?: string;
}) {
return (
<div
data-slot="image-content-filter-error"
className={cn(
"flex min-h-32 flex-col items-center justify-center gap-2 bg-muted/50 p-4 text-center",
className,
)}
>
<ShieldAlertIcon className="size-8 text-muted-foreground" />
<p className="font-medium text-sm">Image could not be generated</p>
{reason && <p className="text-muted-foreground text-xs">{reason}</p>}
</div>
);
}
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<void>;
className?: string;
};
function RegenerateButton({
onRegenerate,
}: {
onRegenerate: () => void | Promise<void>;
}) {
const [isRegenerating, setIsRegenerating] = useState(false);
return (
<button
type="button"
onClick={async () => {
setIsRegenerating(true);
try {
await onRegenerate();
} finally {
setIsRegenerating(false);
}
}}
disabled={isRegenerating}
data-slot="image-regenerate"
aria-label="Regenerate image"
className="inline-flex size-7 items-center justify-center rounded hover:bg-muted disabled:opacity-50"
>
<RefreshCwIcon
className={cn("size-4", isRegenerating && "animate-spin")}
/>
</button>
);
}
function ImageActions({ part, onRegenerate, className }: ImageActionsProps) {
return (
<div
data-slot="image-actions"
className={cn("flex items-center gap-1 p-1", className)}
>
<button
type="button"
onClick={() => downloadImagePart(part)}
data-slot="image-download"
aria-label="Download image"
className="inline-flex size-7 items-center justify-center rounded hover:bg-muted"
>
<DownloadIcon className="size-4" />
</button>
<button
type="button"
onClick={() => {
copyImagePart(part).catch((error) => {
reportImageCopyError(error);
});
}}
data-slot="image-copy"
aria-label="Copy image"
className="inline-flex size-7 items-center justify-center rounded hover:bg-muted"
>
<CopyIcon className="size-4" />
</button>
{onRegenerate && <RegenerateButton onRegenerate={onRegenerate} />}
</div>
);
}
const ImageImpl: ImageMessagePartComponent = (props) => {
const { image, filename, status } = props;
const alt = filename || "Image content";
if (status?.type === "running") {
return (
<ImageRoot>
<ImageGenerating />
<ImageFilename>{filename}</ImageFilename>
</ImageRoot>
);
}
if (status?.type === "incomplete" && status.reason === "content-filter") {
return (
<ImageRoot>
<ImageContentFilterError reason="The provider blocked this image." />
</ImageRoot>
);
}
return (
<ImageRoot>
<ImageZoom src={image} alt={alt}>
<ImagePreview src={image} alt={alt} />
</ImageZoom>
<ImageFilename>{filename}</ImageFilename>
</ImageRoot>
);
};
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,
};

50
studio/frontend/src/components/assistant-ui/message-timing.tsx
View file

@ -33,10 +33,24 @@ export const MessageTiming: FC<{
if (timing?.totalStreamTime === undefined) return null;
const serverTimings = (
const custom = (
message.metadata as Record<string, unknown> | undefined
)?.custom as { serverTimings?: Record<string, number> } | undefined;
const st = serverTimings?.serverTimings;
)?.custom as
| {
serverTimings?: Record<string, number>;
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<{
</span>
</div>
)}
{(st?.cache_n ?? 0) > 0 && (
{cacheHits > 0 && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Cache hits</span>
<span className="font-mono tabular-nums">
{formatNumber(st!.cache_n)}
{formatNumber(cacheHits)}
</span>
</div>
)}
{cacheWrites > 0 && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Cache writes</span>
<span className="font-mono tabular-nums">
{formatNumber(cacheWrites)}
</span>
</div>
)}
@ -146,7 +168,7 @@ export const MessageTiming: FC<{
</>
) : (
<>
{/* Client-side metrics (safetensors fallback) */}
{/* Client-side metrics (safetensors + external provider fallback) */}
{timing.firstTokenTime !== undefined && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">First token</span>
@ -155,6 +177,22 @@ export const MessageTiming: FC<{
</span>
</div>
)}
{cacheHits > 0 && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Cache hits</span>
<span className="font-mono tabular-nums">
{formatNumber(cacheHits)}
</span>
</div>
)}
{cacheWrites > 0 && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Cache writes</span>
<span className="font-mono tabular-nums">
{formatNumber(cacheWrites)}
</span>
</div>
)}
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Total</span>
<span className="font-mono tabular-nums">

2
studio/frontend/src/components/assistant-ui/reasoning.tsx
View file

@ -134,7 +134,7 @@ function ReasoningTrigger({
{active ? (
<span className="text-sm">Thinking...</span>
) : (
<span>Thought for {duration ?? 0} seconds</span>
<span>Thought for {duration ?? 0} {duration === 1 ? "second" : "seconds"}</span>
)}
</span>
<ChevronDownIcon

18
studio/frontend/src/components/assistant-ui/sources.tsx
View file

@ -127,6 +127,12 @@ function Source({
// ── Source badge with hover card ─────────────────────────────
interface SourceData {
/**
* Stable per-citation key. Two Anthropic document citations into
* different spans of the same source share a ``url``, so React keys
* on ``id`` to keep each footnote distinct.
*/
id: string;
url: string;
title: string;
description?: string;
@ -190,8 +196,14 @@ const SourcesGroup: FC = () => {
"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({
url: part.url as string,
id: partId,
url,
title: (part as { title?: string }).title || "",
description: (part as { metadata?: { description?: string } })
.metadata?.description,
@ -258,7 +270,7 @@ const SourcesGroup: FC = () => {
className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none"
>
{sources.map((source) => (
<span key={source.url} className="inline-block">
<span key={source.id} className="inline-block">
<Source href={source.url}>
<SourceIcon url={source.url} />
<SourceTitle>{source.title || extractDomain(source.url)}</SourceTitle>
@ -270,7 +282,7 @@ const SourcesGroup: FC = () => {
{/* Visible container */}
<div className="flex flex-wrap gap-1">
{displayedSources.map((source) => (
<SourceBadge key={source.url} source={source} />
<SourceBadge key={source.id} source={source} />
))}
{shouldCollapse && !expanded && (
<button

486
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -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";
@ -39,13 +44,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,
@ -69,7 +75,6 @@ import {
DownloadIcon,
GlobeIcon,
HeadphonesIcon,
ImageIcon,
LightbulbIcon,
LightbulbOffIcon,
MicIcon,
@ -79,30 +84,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
@ -113,85 +119,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 (
<ThreadPrimitive.Root
className="aui-root aui-thread-root @container relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
style={{
["--thread-max-width" as string]: "48rem",
["--thread-content-max-width" as string]:
"calc(var(--thread-max-width) - 1.5rem)",
}}
>
<IntentAwareScrollProvider value={autoScrollContext}>
<ThreadPrimitive.Viewport
ref={viewportRef}
autoScroll={false}
scrollToBottomOnRunStart={false}
scrollToBottomOnInitialize={false}
scrollToBottomOnThreadSwitch={false}
className={cn(
"aui-thread-viewport aui-stream-viewport relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-x-auto overflow-y-auto scroll-smooth px-5",
hideComposer ? "pt-4" : "pt-[48px]",
)}
>
{!hideWelcome && (
<AuiIf condition={({ thread }) => thread.isEmpty && !thread.isLoading}>
<ThreadWelcome hideComposer={hideComposer} />
</AuiIf>
)}
<GeneratedImageOverlayProvider key={threadId ?? "default"} threadId={threadId}>
<ThreadPrimitive.Root
className="aui-root aui-thread-root @container relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
style={{
["--thread-max-width" as string]: "48rem",
["--thread-content-max-width" as string]:
"calc(var(--thread-max-width) - 1.5rem)",
}}
>
<IntentAwareScrollProvider value={autoScrollContext}>
<ThreadPrimitive.Viewport
ref={viewportRef}
autoScroll={false}
scrollToBottomOnRunStart={false}
scrollToBottomOnInitialize={false}
scrollToBottomOnThreadSwitch={false}
className={cn(
"aui-thread-viewport aui-stream-viewport relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-x-auto overflow-y-auto scroll-smooth px-5",
hideComposer ? "pt-4" : "pt-[48px]",
)}
>
{!hideWelcome && (
<AuiIf
condition={({ thread }) => thread.isEmpty && !thread.isLoading}
>
<ThreadWelcome hideComposer={hideComposer} threadId={threadId} />
</AuiIf>
)}
<ThreadPrimitive.Messages
components={{
UserMessage,
EditComposer,
AssistantMessage,
}}
/>
<ThreadPrimitive.Messages
components={{
UserMessage,
EditComposer,
AssistantMessage,
}}
/>
{/* 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. */}
<AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
<div
className={cn("shrink-0", hideComposer ? "h-16" : "h-40")}
aria-hidden={true}
/>
</AuiIf>
<AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
<ThreadPrimitive.ViewportFooter
className={cn(
"aui-thread-viewport-footer pointer-events-none sticky z-20 flex w-full justify-center bg-transparent",
hideComposer ? "bottom-3" : "bottom-[140px]",
)}
>
<ThreadScrollToBottom />
</ThreadPrimitive.ViewportFooter>
</AuiIf>
</ThreadPrimitive.Viewport>
{!hideComposer && (
<AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
<div className="aui-thread-composer-dock pointer-events-none absolute bottom-0 left-0 right-0 md:right-[10px] z-20">
<AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
<div
className={cn("shrink-0", hideComposer ? "h-16" : "h-40")}
aria-hidden={true}
className="absolute inset-x-0 bottom-0 top-[10px] bg-background"
/>
<div className="relative px-5 pb-2">
<div className="pointer-events-auto mx-auto w-full max-w-(--thread-max-width)">
<ComposerAnimated disabled={isComposerAttachPending} />
</div>
<p className="composer-footer-note">
LLMs can make mistakes. Double-check responses.
</p>
</div>
</div>
</AuiIf>
</AuiIf>
<AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
<ThreadPrimitive.ViewportFooter
className={cn(
"aui-thread-viewport-footer pointer-events-none sticky z-20 flex w-full justify-center bg-transparent",
// 150px (was 140px) to add a small gap above the composer
hideComposer ? "bottom-3" : "bottom-[150px]",
)}
>
<ThreadScrollToBottom />
</ThreadPrimitive.ViewportFooter>
</AuiIf>
</ThreadPrimitive.Viewport>
<GeneratedImageViewportOverlay hideComposer={hideComposer} />
{!hideComposer && (
<AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
<ThreadComposerDock
disabled={isComposerAttachPending}
threadId={threadId}
/>
</AuiIf>
)}
</IntentAwareScrollProvider>
</ThreadPrimitive.Root>
</GeneratedImageOverlayProvider>
);
};
const GeneratedImageViewportOverlay: FC<{ hideComposer?: boolean }> = ({
hideComposer,
}) => {
const { overlay, closeOverlay } = useGeneratedImageOverlay();
useEffect(() => {
if (!overlay) {
return;
}
document.querySelector<HTMLTextAreaElement>(".aui-composer-input")?.focus();
}, [overlay]);
if (!overlay) {
return null;
}
return (
<div className="pointer-events-none absolute inset-0 z-30">
<button
type="button"
className="pointer-events-auto absolute inset-0 bg-background/65 backdrop-blur-[1px] dark:bg-background/55"
onClick={closeOverlay}
aria-label="Close generated image preview"
/>
<section
className={cn(
"pointer-events-none absolute inset-x-5 top-[48px] flex flex-col items-center",
hideComposer ? "bottom-4" : "bottom-[150px]",
)}
</IntentAwareScrollProvider>
</ThreadPrimitive.Root>
aria-label="Generated image preview"
>
<div className="pointer-events-auto relative flex min-h-0 w-full max-w-[1100px] flex-1 flex-col items-center justify-center gap-3 rounded-3xl bg-muted/10 p-3 ring-1 ring-border/20">
<div className="absolute inset-x-3 top-3 z-10 flex justify-end">
<div className="flex shrink-0 items-center gap-1 rounded-full bg-background/70 p-1 ring-1 ring-border/20 backdrop-blur-sm">
<Button
type="button"
variant="ghost"
size="icon-sm"
className="size-7 rounded-full"
onClick={() =>
downloadImagePart({
image: overlay.image,
filename: overlay.filename,
})
}
aria-label="Download generated image"
>
<DownloadIcon className="size-3.5" />
</Button>
<Button
type="button"
variant="ghost"
size="icon-sm"
className="size-7 rounded-full"
onClick={closeOverlay}
aria-label="Close generated image preview"
>
<XIcon className="size-3.5" />
</Button>
</div>
</div>
<div className="flex min-h-0 flex-1 items-center justify-center pt-1">
<img
src={overlay.image}
alt={overlay.title}
className="max-h-full max-w-full object-contain"
/>
</div>
<div
className="w-full max-w-[min(100%,46rem)] shrink-0 text-center"
title={overlay.title}
>
<p className="truncate text-xs font-semibold text-foreground/80">
Generated image
</p>
{overlay.metadata ? (
<p className="truncate text-[11px] font-medium text-muted-foreground">
{overlay.metadata}
</p>
) : null}
{hideComposer ? null : (
<p className="mx-auto mt-2 inline-flex rounded-full bg-primary/10 px-3 py-1 text-xs font-medium text-primary">
Type edits below, then send
</p>
)}
</div>
</div>
</section>
</div>
);
};
const ThreadComposerDock: FC<{
disabled?: boolean;
threadId?: string | null;
}> = ({ disabled, threadId }) => {
const { overlay } = useGeneratedImageOverlay();
return (
<div
className={cn(
"aui-thread-composer-dock pointer-events-none absolute bottom-0 left-0 right-0 md:right-[10px]",
overlay ? "z-40" : "z-20",
)}
>
<div
aria-hidden={true}
className="absolute inset-x-0 bottom-0 top-[10px] bg-background"
/>
<div className="relative px-5 pb-2">
<div className="pointer-events-auto mx-auto w-full max-w-(--thread-max-width)">
<ComposerAnimated disabled={disabled} threadId={threadId} />
</div>
<p className="composer-footer-note">
LLMs can make mistakes. Double-check responses.
</p>
</div>
</div>
);
};
@ -219,13 +345,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");
}, []);
@ -240,11 +370,7 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
<div className="aui-thread-welcome-center flex w-full grow flex-col items-center justify-center pb-[48px]">
<div className="aui-thread-welcome-message flex w-full flex-col justify-center gap-6 px-4">
<div className="flex flex-col items-center gap-2 text-center">
<img
src={currentEmojiSrc}
alt="Sloth mascot"
className="size-20"
/>
<img src={currentEmojiSrc} alt="Sloth mascot" className="size-20" />
<h1 className="aui-thread-welcome-message-inner fade-in slide-in-from-bottom-1 animate-in font-heading font-semibold text-2xl tracking-[-0.02em] duration-200">
Chat with your model
</h1>
@ -252,18 +378,21 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
Run GGUFs, safetensors, vision and audio models
</p>
</div>
{!hideComposer && <ComposerAnimated />}
{!hideComposer && <ComposerAnimated threadId={threadId} />}
</div>
</div>
</div>
);
};
const ComposerAnimated: FC<{ disabled?: boolean }> = ({ disabled }) => {
const ComposerAnimated: FC<{
disabled?: boolean;
threadId?: string | null;
}> = ({ disabled, threadId }) => {
return (
<div className="relative mx-auto min-w-0 w-full max-w-(--thread-max-width)">
<div className="relative z-10 w-full">
<Composer disabled={disabled} />
<Composer disabled={disabled} threadId={threadId} />
</div>
</div>
);
@ -293,8 +422,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,
@ -304,22 +446,78 @@ 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 referenceThreadId = threadId ?? activeThreadId ?? null;
const hasSendableContent =
composerText.trim().length > 0 || hasAttachments || hasPendingAudio;
const shouldBlockSend = useCallback(
() =>
!hasSendableContent || isComposingRef.current || hasPendingAttachments,
[hasPendingAttachments, hasSendableContent, isComposingRef],
);
const handleSubmit = useCallback(
(event: FormEvent<HTMLFormElement>) => {
if (
disabled ||
!hasSendableContent ||
isComposingRef.current ||
hasPendingAttachments
) {
(event: Parameters<NonNullable<ComponentProps<"form">["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();
}
},
[disabled, hasPendingAttachments, hasSendableContent, isComposingRef],
[
aui,
closeOverlay,
composerText,
disabled,
overlay,
referenceThreadId,
setImageToolsEnabled,
setPendingImageEditReference,
shouldBlockSend,
],
);
const composerContent = (
@ -328,13 +526,15 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
<PendingAudioChip />
<ToolStatusDisplay />
<ComposerPrimitive.Input
placeholder="Send a message..."
placeholder={
overlay ? "Type your edits for your image" : "Send a message..."
}
className="aui-composer-input composer-input"
minRows={1}
maxRows={12}
autoFocus={!disabled}
disabled={disabled}
aria-label="Message input"
aria-label={overlay ? "Image edit instructions" : "Message input"}
// dir="auto": browser picks LTR/RTL from the first strong char;
// no effect on Latin / CJK / Devanagari.
dir="auto"
@ -342,11 +542,12 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
/>
<ComposerAction
disabled={
disabled || !hasSendableContent || isComposing || hasPendingAttachments
}
blockSend={() =>
!hasSendableContent || isComposingRef.current || hasPendingAttachments
disabled ||
!hasSendableContent ||
isComposing ||
hasPendingAttachments
}
shouldBlockSend={shouldBlockSend}
/>
</>
);
@ -553,7 +754,6 @@ const ComposerAudioUpload: FC = () => {
);
};
const ReasoningToggle: FC = () => {
const modelLoaded = useChatRuntimeStore(
(s) => !!s.params.checkpoint && !s.modelLoading,
@ -565,8 +765,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,
@ -619,7 +823,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") ||
@ -677,23 +882,25 @@ const ReasoningToggle: FC = () => {
{effectiveReasoningEffortLevels
.filter((level) => level !== "none")
.map((level) => (
<DropdownMenuItem
key={level}
onSelect={() => {
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" : ""}
</DropdownMenuItem>
))}
<DropdownMenuItem
key={level}
onSelect={() => {
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"
: ""}
</DropdownMenuItem>
))}
</DropdownMenuContent>
</DropdownMenu>
);
@ -808,8 +1015,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 (
<button
@ -899,10 +1105,12 @@ const ImagesToggle: FC = () => {
className="composer-pill-btn"
data-active={imageToolsEnabled && !disabled ? "true" : "false"}
aria-label={
imageToolsEnabled ? "Disable image generation" : "Enable image generation"
imageToolsEnabled
? "Disable image generation"
: "Enable image generation"
}
>
<ImageIcon className="size-3.5" />
<HugeiconsIcon icon={Image03Icon} className="size-3.5" strokeWidth={2} />
<span>Images</span>
</button>
);
@ -966,10 +1174,10 @@ const ToolStatusDisplay: FC = () => {
);
};
const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
disabled,
blockSend,
}) => {
const ComposerAction: FC<{
disabled?: boolean;
shouldBlockSend?: () => boolean;
}> = ({ disabled, shouldBlockSend }) => {
return (
<div className="aui-composer-action-wrapper composer-action-wrapper">
<div className="flex items-center gap-0.5">
@ -1016,7 +1224,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
size="icon"
disabled={disabled}
onClick={(event) => {
if (blockSend?.()) {
if (shouldBlockSend?.()) {
event.preventDefault();
}
}}
@ -1284,7 +1492,11 @@ const UserActionBar: FC = () => {
<CopyButton />
<ActionBarPrimitive.Edit asChild={true}>
<TooltipIconButton tooltip="Edit" className="aui-user-action-edit">
<HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
<HugeiconsIcon
icon={Edit03Icon}
strokeWidth={1.75}
className="size-icon"
/>
</TooltipIconButton>
</ActionBarPrimitive.Edit>
<DeleteMessageButton />

322
studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx
View file

@ -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 (
<span
key={index}
className="generated-image-loading-dot"
style={
{
"--dot-row": row,
"--dot-col": col,
} as CSSProperties
}
/>
);
});
function GeneratedImagePlaceholder({ label }: { label: string }) {
return (
<div
className={cn(
"generated-image-loading-card flex aspect-square w-[480px] max-w-full items-center justify-center rounded-2xl bg-muted/20 shadow-[0_0_12px_rgba(15,23,42,0.05),0_6px_18px_rgba(15,23,42,0.04)] dark:shadow-[0_0_12px_rgba(0,0,0,0.18),0_6px_18px_rgba(0,0,0,0.12)]",
)}
aria-busy="true"
aria-label={label}
aria-live="polite"
>
<span className="sr-only">{label}</span>
<div className="generated-image-loading-wave" aria-hidden={true}>
{loadingDots}
</div>
</div>
);
}
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<HTMLDivElement | null>(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<HTMLButtonElement>,
) => {
event.preventDefault();
event.stopPropagation();
};
const handleDownload = (event: MouseEvent<HTMLButtonElement>) => {
stopOverlayActionPropagation(event);
if (imagePart) {
downloadImagePart(imagePart);
}
};
const handleEditClick = (event: MouseEvent<HTMLButtonElement>) => {
stopOverlayActionPropagation(event);
showPreview();
};
if (isPendingImage) {
return (
<div className="aui-tool-fallback-root w-full py-1">
<GeneratedImagePlaceholder label={runningLabel} />
</div>
);
}
return (
<ToolFallbackRoot open={open} onOpenChange={setOpen}>
@ -102,21 +288,77 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
icon={ImageIcon}
/>
<ToolFallbackContent>
{isRunning && !imageSrc ? (
<div className="flex items-center gap-2 text-sm text-muted-foreground">
<LoaderIcon className="size-3.5 animate-spin" />
<span>{runningLabel}</span>
</div>
) : imageSrc ? (
<figure className="m-0 flex flex-col gap-1.5">
<img
src={imageSrc}
alt={prompt || "Generated image"}
className="max-w-full rounded-md border border-border/60"
/>
{prompt ? (
<figcaption className="text-xs leading-snug text-muted-foreground">
{prompt}
{imagePart ? (
<figure className="m-0 flex flex-col gap-2">
<div className="group/generated-image relative aspect-square w-[480px] max-w-full overflow-hidden rounded-2xl bg-muted/25 shadow-lg shadow-foreground/5 dark:shadow-black/25">
<img
src={imagePart.image}
alt=""
aria-hidden={true}
className="pointer-events-none absolute inset-0 size-full scale-110 object-cover opacity-25 blur-2xl saturate-125"
/>
<div className="pointer-events-none absolute inset-0 bg-background/45" />
<button
type="button"
className="relative z-10 block size-full cursor-zoom-in overflow-hidden rounded-2xl focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
onClick={showPreview}
aria-label="Open generated image preview"
>
<Image.Preview
src={imagePart.image}
alt={imageTitle}
containerClassName="flex size-full min-h-0 items-center justify-center bg-transparent"
className="size-full object-contain"
/>
</button>
<div className="pointer-events-none absolute inset-x-0 bottom-0 z-20 flex items-end justify-between gap-2 bg-gradient-to-t from-black/55 via-black/20 to-transparent p-3 opacity-100 transition-opacity sm:opacity-0 sm:group-hover/generated-image:opacity-100 sm:group-focus-within/generated-image:opacity-100">
<Button
type="button"
variant="dark"
size="sm"
className="pointer-events-auto h-8 rounded-full bg-black/70 text-white hover:bg-black/85"
onClick={handleEditClick}
>
<PencilIcon className="size-3.5" />
Edit
</Button>
<Button
type="button"
variant="dark"
size="icon-sm"
className="pointer-events-auto rounded-full bg-black/70 text-white hover:bg-black/85"
onClick={handleDownload}
aria-label="Download generated image"
>
<DownloadIcon className="size-4" />
</Button>
</div>
</div>
{captionPrompt ? (
<figcaption className="max-w-[480px] text-xs leading-5 text-muted-foreground">
<div
ref={captionRef}
className={cn(
"whitespace-pre-wrap break-words",
shouldClampPrompt && "max-h-20 overflow-hidden",
)}
>
{captionPrompt}
</div>
{promptCanExpand ? (
<button
type="button"
className="mt-2 inline-flex text-xs font-medium text-foreground/80 underline-offset-4 hover:text-foreground hover:underline focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
onClick={() =>
setExpandedCaptionPrompt((value) =>
value === captionPrompt ? null : captionPrompt,
)
}
aria-expanded={promptExpanded}
>
{promptExpanded ? "Show less" : "Show more"}
</button>
) : null}
</figcaption>
) : null}
</figure>

31
studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx
View file

@ -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 <a href>.
*/
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;
}

596
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -1,7 +1,7 @@
// 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 { getAuthToken } from "@/features/auth/session";
import { getAuthToken } from "@/features/auth";
import { apiUrl } from "@/lib/api-base";
import { toast } from "@/lib/toast";
import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
@ -21,6 +21,7 @@ import { pickFriendlyContainerName } from "../lib/friendly-names";
import {
EXTERNAL_MAX_OUTPUT_TOKENS,
clampReasoningEffortToLevels,
getExternalMaxOutputTokens,
getExternalMinOutputTokens,
getExternalReasoningCapabilities,
getProviderCapabilities,
@ -28,13 +29,19 @@ import {
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 { getImageInputUnavailableReason } from "../utils/image-input-support";
@ -70,6 +77,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. */
@ -143,6 +157,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 <a href>.
*/
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<string, unknown>,
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";
@ -167,7 +266,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 <a href>.
const url = isSafeNavigableSourceUrl(urlMatch[1]);
if (!url) continue;
const snippet = snippetMatch?.[1]?.trim();
sources.push({
type: "source" as const,
@ -302,37 +405,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);
}
}
}
@ -340,6 +436,78 @@ function collectImageParts(
return parts;
}
function normalizeOpenAIReasoningItem(
value: unknown,
): OpenAIReasoningContentPart | null {
if (!value || typeof value !== "object") {
return null;
}
const item = value as Record<string, unknown>;
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<string, unknown>;
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<OpenAIMessageContent, string> = [];
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<string, unknown> }
| undefined;
return metadata?.custom?.anthropicRefusal === true;
}
function toOpenAIMessage(message: RunMessage): {
role: "system" | "user" | "assistant";
content: OpenAIMessageContent;
@ -360,16 +528,27 @@ function toOpenAIMessage(message: RunMessage): {
/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 null;
}
}
const imageParts = collectImageParts(message);
if (imageParts.length > 0) {
return {
role: message.role,
content: [{ type: "text", text: textContent }, ...imageParts],
content: [
...(textContent ? [{ type: "text" as const, text: textContent }] : []),
...imageParts,
],
};
}
if (!textContent) {
return null;
}
return { role: message.role, content: textContent };
}
@ -804,17 +983,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
@ -826,6 +1040,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
: "Pick a model in the top bar, then retry.",
},
);
clearSelectedImageEditReference();
throw new Error("Load a model first.");
}
}
@ -833,7 +1048,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// Re-read store after potential auto-load / model ready wait
runtime = useChatRuntimeStore.getState();
const { params } = runtime;
const { supportsTools, toolsEnabled, codeToolsEnabled, imageToolsEnabled } = runtime;
const {
supportsTools,
toolsEnabled,
codeToolsEnabled,
imageToolsEnabled,
webFetchToolsEnabled,
} = runtime;
const externalSelection = parseExternalModelId(params.checkpoint);
const isExternalRequest = externalSelection !== null;
if (
@ -844,6 +1065,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
@ -859,6 +1081,7 @@ 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.
@ -869,36 +1092,34 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
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 webSearchEnabledForThisTurn = Boolean(
externalProvider &&
toolsEnabled &&
providerSupportsBuiltinWebSearch(externalProvider.providerType),
);
const codeExecEnabledForThisTurn = Boolean(
externalProvider &&
externalSelection &&
codeToolsEnabled &&
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),
@ -918,11 +1139,58 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
),
);
const outboundMessages = messages
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.");
}
// Two-pass build: a refused assistant turn also drops the user
// prompt that triggered it (leaving it in context re-triggers
// the classifier). Refusal flag rides assistant
// metadata.custom.anthropicRefusal, set out-of-band from the
// backend _toolEvent.
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);
}
const outboundMessages = survivingMessages
.map(toOpenAIMessage)
.filter((message): message is NonNullable<typeof message> =>
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;
}
}
outboundMessages.splice(insertAt, 0, referenceMessage);
}
const safeSystemPrompt =
typeof params.systemPrompt === "string" ? params.systemPrompt : "";
@ -941,24 +1209,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.";
@ -988,8 +1283,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
@ -1018,6 +1315,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
const gatedThreadKey = resolvedThreadId || "__default";
runtime.setThreadRunning(gatedThreadKey, true);
runtime.setThreadRunning(gatedThreadKey, false);
clearSelectedImageEditReference();
throw new Error(imageGateReason);
}
}
@ -1025,7 +1323,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);
@ -1136,6 +1434,32 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// Tool call content parts — accumulated and yielded cumulatively.
// result is set directly on the tool-call part when tool_end arrives.
const toolCallParts: ToolCallMessagePart[] = [];
const orderAssistantContent = (
textParts: ReturnType<typeof parseAssistantContent>,
) => {
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;
@ -1314,8 +1638,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
) {
void updateStoredChatThreadEventually(t.id, {
openaiCodeExecContainerId: null,
})
.catch(() => {});
}).catch(() => {});
continue;
}
openaiCodeExecContainerId = t.openaiCodeExecContainerId;
@ -1359,8 +1682,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
openaiCodeExecContainerId = created.id;
void updateStoredChatThreadEventually(resolvedThreadId, {
openaiCodeExecContainerId: created.id,
})
.catch(() => {});
}).catch(() => {});
} catch {
// Fall back to backend's container_auto path on
// failure — keeps the chat moving; the next turn
@ -1382,18 +1704,17 @@ 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
@ -1419,13 +1740,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.
// web_fetch has its own Fetch pill, independent
// of Search. Anthropic-only today.
...(webFetchEnabledForThisTurn ? ["web_fetch"] : []),
...(codeExecEnabledForThisTurn ? ["code_execution"] : []),
// OpenAI Responses-API only: `image_generation`
@ -1473,11 +1789,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// attaches `cache_control.ttl` when the value is one of
// "5m" / "1h" (see external_provider.py near line 1375),
// so unknown values are a no-op end-to-end.
...(supportsProviderPromptCacheTtl(externalProvider.providerType) &&
...(supportsProviderPromptCacheTtl(
externalProvider.providerType,
) &&
(externalProvider.enablePromptCaching ?? true) &&
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
@ -1541,10 +1869,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
@ -1583,6 +1916,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<string, unknown>,
idx,
);
if (
part &&
!documentCitationParts.some((p) => p.id === part.id)
) {
documentCitationParts.push(part);
}
});
}
continue;
}
if (toolEvent.type === "container_invalidated") {
if (resolvedThreadId) {
const field =
@ -1591,11 +1945,16 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
: "openaiCodeExecContainerId";
void updateStoredChatThreadEventually(resolvedThreadId, {
[field]: null,
})
.catch(() => {});
}).catch(() => {});
}
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) ||
@ -1630,6 +1989,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
size?: string;
quality?: string;
background?: string;
prompt?: string;
};
const imageB64 = toolEvent.image_b64 as string | undefined;
if (
@ -1651,6 +2011,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);
@ -1668,16 +2029,29 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
} else {
parsedResult = rawResult;
}
const nextArgs =
toolEvent.arguments &&
typeof toolEvent.arguments === "object"
? (toolEvent.arguments as ToolCallMessagePart["args"])
: undefined;
const mergedArgs = nextArgs
? { ...(toolCallParts[idx].args ?? {}), ...nextArgs }
: toolCallParts[idx].args;
toolCallParts[idx] = {
...toolCallParts[idx],
args: mergedArgs,
argsText: mergedArgs
? JSON.stringify(mergedArgs)
: toolCallParts[idx].argsText,
result: parsedResult,
};
}
}
// Yield cumulative state so tool UI updates (tools first, text after)
// Yield cumulative state so tool UI updates. Search/code tools stay
// before the text, while generated images sit after the answer.
const textParts = parseAssistantContent(cumulativeText);
yield {
content: [...toolCallParts, ...textParts],
content: orderAssistantContent(textParts),
metadata: {
timing: buildTiming(
streamStartTime,
@ -1825,7 +2199,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
if (parts.length > 0 || toolCallParts.length > 0) {
yield {
content: [...toolCallParts, ...parts],
content: orderAssistantContent(parts),
metadata: {
timing: buildTiming(
streamStartTime,
@ -1881,18 +2255,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,
});
}
@ -1908,21 +2295,24 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
yield {
content: [
...toolCallParts,
...parseAssistantContent(cumulativeText),
...orderAssistantContent(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,

59
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -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";
@ -779,6 +781,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 +806,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 +842,7 @@ export function ChatPage(): ReactElement {
supportsBuiltinWebSearch,
supportsBuiltinCodeExecution,
supportsBuiltinImageGeneration,
supportsBuiltinWebFetch,
toolsEnabled: nextToolsEnabled,
codeToolsEnabled: supportsBuiltinCodeExecution
? (storedCodeToolsEnabled ?? false)
@ -841,6 +850,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(() => {
@ -1008,6 +1021,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 +1042,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 +1056,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 +1086,7 @@ export function ChatPage(): ReactElement {
supportsBuiltinWebSearch,
supportsBuiltinCodeExecution,
supportsBuiltinImageGeneration,
supportsBuiltinWebFetch,
toolsEnabled: nextToolsEnabled,
codeToolsEnabled: supportsBuiltinCodeExecution
? (storedCodeToolsEnabled ?? false)
@ -1070,6 +1094,9 @@ export function ChatPage(): ReactElement {
imageToolsEnabled: supportsBuiltinImageGeneration
? (storedImageToolsEnabled ?? false)
: false,
webFetchToolsEnabled: supportsBuiltinWebFetch
? (storedWebFetchToolsEnabled ?? false)
: false,
...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }),
});
return;
@ -1161,7 +1188,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 +1204,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 +1542,13 @@ export function ChatPage(): ReactElement {
) : null}
</div>
<div className="ml-auto flex items-center gap-2">
{view.mode === "single" && ggufContextLength && contextUsage ? (
{view.mode === "single" && contextUsage ? (
<ContextUsageBar
used={contextUsage.totalTokens}
// null on external providers; the bar handles that.
total={ggufContextLength}
cached={contextUsage.cachedTokens}
cacheWrites={contextUsage.cacheWriteTokens}
promptTokens={contextUsage.promptTokens}
completionTokens={contextUsage.completionTokens}
className="h-[34px]"

35
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -85,8 +85,10 @@ import {
import {
EXTERNAL_MAX_OUTPUT_TOKENS,
type ProviderCapabilities,
getExternalMaxOutputTokens,
getExternalMinOutputTokens,
providerSupportsBuiltinCodeExecution,
providerSupportsFastMode,
} from "./provider-capabilities";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import type { InferenceParams } from "./types/runtime";
@ -552,6 +554,12 @@ export function ChatSettingsPanel({
activeExternalProvider.baseUrl,
) &&
activeExternalProvider.providerType === "openai";
const showFastModeControl =
activeExternalProvider != null &&
providerSupportsFastMode(
activeExternalProvider.providerType,
externalSelection?.modelId,
);
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const openAiApiKeyForSection = activeExternalProvider
? getExternalProviderApiKey(activeExternalProvider.id) || null
@ -1152,6 +1160,28 @@ export function ChatSettingsPanel({
</Select>
</div>
) : null}
{showFastModeControl ? (
<div className="flex items-center justify-between gap-3 pt-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Fast mode
</span>
<InfoHint>
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.
</InfoHint>
</div>
<Switch
className="panel-switch shrink-0"
checked={Boolean(params.fastMode)}
onCheckedChange={set("fastMode")}
aria-label="Fast mode"
/>
</div>
) : null}
</CollapsibleSection>
) : null}
@ -1280,7 +1310,10 @@ export function ChatSettingsPanel({
}
max={
isExternalModel
? EXTERNAL_MAX_OUTPUT_TOKENS
? getExternalMaxOutputTokens(
externalProviderType,
externalSelection?.modelId,
)
: isGguf && ggufContextLength
? ggufContextLength
: 32768

89
studio/frontend/src/features/chat/components/context-usage-bar.tsx
View file

@ -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 (
<Tooltip>
<TooltipTrigger asChild>
<button
type="button"
aria-label={`Context usage: ${formatTokenCount(used)} of ${formatTokenCount(total)} tokens`}
aria-label={
hasKnownLimit
? `Context usage: ${formatTokenCount(used)} of ${formatTokenCount(total as number)} tokens`
: `Token usage: ${formatTokenCount(used)} tokens`
}
className={cn(
"flex items-center gap-2 rounded-[10px] px-2.5 py-1 font-mono text-chat-icon-fg text-[13px] tabular-nums transition-colors hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover",
className,
)}
>
<span>
{formatTokenCount(used)} / {formatTokenCount(total)}
{hasKnownLimit
? `${formatTokenCount(used)} / ${formatTokenCount(total as number)}`
: `${formatTokenCount(used)} tokens`}
</span>
<div className="h-1.5 w-16 rounded-full bg-black/10 dark:bg-white/15 overflow-hidden">
<div
className={cn("h-full rounded-full transition-all", severity.bar)}
style={{ width: `${percent}%` }}
/>
</div>
{hasKnownLimit && percent !== null ? (
<div className="h-1.5 w-16 rounded-full bg-black/10 dark:bg-white/15 overflow-hidden">
<div
className={cn("h-full rounded-full transition-all", severity.bar)}
style={{ width: `${percent}%` }}
/>
</div>
) : null}
</button>
</TooltipTrigger>
<TooltipContent
@ -68,12 +97,14 @@ export const ContextUsageBar: FC<{
className="[&_span>svg]:hidden!"
>
<div className="grid min-w-44 gap-1.5 text-xs">
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Context usage</span>
<span className={cn("font-mono tabular-nums font-medium", severity.text)}>
{percent.toFixed(1)}%
</span>
</div>
{hasKnownLimit && percent !== null ? (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Context usage</span>
<span className={cn("font-mono tabular-nums font-medium", severity.text)}>
{percent.toFixed(1)}%
</span>
</div>
) : null}
{promptTokens !== undefined && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Prompt tokens</span>
@ -98,20 +129,32 @@ export const ContextUsageBar: FC<{
</span>
</div>
)}
{cacheWrites !== undefined && cacheWrites > 0 && (
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Cache writes</span>
<span className="font-mono tabular-nums">
{formatTokenCountFull(cacheWrites)}
</span>
</div>
)}
<div className="my-0.5 border-t border-border/40" />
<div className="flex items-center justify-between gap-4">
<span className="text-muted-foreground">Total</span>
<span className="text-muted-foreground">
{hasKnownLimit ? "Total" : "Total tokens"}
</span>
<span className="font-mono tabular-nums">
{formatTokenCountFull(used)} / {formatTokenCountFull(total)}
{hasKnownLimit
? `${formatTokenCountFull(used)} / ${formatTokenCountFull(total as number)}`
: formatTokenCountFull(used)}
</span>
</div>
{percent > 85 && (
{hasKnownLimit && percent !== null && percent > 85 ? (
<div className="mt-1 max-w-64 text-[11px] leading-snug text-muted-foreground/90">
Close to the context limit. Generation will stop at 100%.
Increase <span className="font-medium">Context Length</span> in
the chat Settings panel to keep going.
</div>
)}
) : null}
</div>
</TooltipContent>
</Tooltip>

8
studio/frontend/src/features/chat/index.ts
View file

@ -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,

2
studio/frontend/src/features/chat/presets/preset-policy.ts
View file

@ -248,7 +248,7 @@ interface BackendInferenceDefaults {
export interface BackendInferenceEnvelope {
is_gguf?: boolean;
context_length?: number | null;
inference?: BackendInferenceDefaults;
inference?: BackendInferenceDefaults | null;
}
export function mergeBackendRecommendedInference({

145
studio/frontend/src/features/chat/provider-capabilities.ts
View file

@ -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
@ -123,11 +200,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 +210,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 +284,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(

21
studio/frontend/src/features/chat/runtime-provider.tsx
View file

@ -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);
}

51
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -21,7 +21,20 @@ 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, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
import {
ArrowUpIcon,
DownloadIcon,
GlobeIcon,
HeadphonesIcon,
LightbulbIcon,
LightbulbOffIcon,
MicIcon,
PlusIcon,
SquareIcon,
XIcon,
} from "lucide-react";
import { Image03Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { toast } from "@/lib/toast";
import { loadModel, validateModel } from "./api/chat-api";
import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers";
@ -34,6 +47,7 @@ import {
getExternalReasoningCapabilities,
providerSupportsBuiltinCodeExecution,
providerSupportsBuiltinImageGeneration,
providerSupportsBuiltinWebFetch,
} from "./provider-capabilities";
import {
type CompositionEvent,
@ -336,6 +350,12 @@ export function SharedComposer({
const setImageToolsEnabled = useChatRuntimeStore(
(s) => s.setImageToolsEnabled,
);
const webFetchToolsEnabled = useChatRuntimeStore(
(s) => s.webFetchToolsEnabled,
);
const setWebFetchToolsEnabled = useChatRuntimeStore(
(s) => s.setWebFetchToolsEnabled,
);
const lastOpenRouterChosenModel = useChatRuntimeStore(
(s) => s.lastOpenRouterChosenModel,
);
@ -426,6 +446,9 @@ export function SharedComposer({
effectiveExternalModelId,
selectedExternalProvider?.baseUrl,
);
const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch(
selectedExternalProvider?.providerType,
);
const searchDisabled =
!modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
const codeDisabled =
@ -437,6 +460,9 @@ export function SharedComposer({
// the pill row stays compact for providers without the capability.
const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration;
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;
@ -1102,10 +1128,31 @@ export function SharedComposer({
imageToolsEnabled ? "Disable image generation" : "Enable image generation"
}
>
<ImageIcon className="size-3.5" />
<HugeiconsIcon
icon={Image03Icon}
className="size-3.5"
strokeWidth={2}
/>
<span>Images</span>
</button>
)}
{showWebFetchPill && (
<button
type="button"
disabled={webFetchDisabled}
onClick={() => setWebFetchToolsEnabled(!webFetchToolsEnabled)}
className="composer-pill-btn"
data-active={
webFetchToolsEnabled && !webFetchDisabled ? "true" : "false"
}
aria-label={
webFetchToolsEnabled ? "Disable URL fetch" : "Enable URL fetch"
}
>
<DownloadIcon className="size-3.5" />
<span>Fetch</span>
</button>
)}
</div>
<div className="flex items-center gap-1">
{dictationSupported && (

82
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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,
@ -25,6 +27,8 @@ export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
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_WEB_FETCH_TOOLS_ENABLED_KEY =
"unsloth_chat_web_fetch_tools_enabled";
// External provider selection is encoded into `params.checkpoint` as
// `external::<providerId>::<modelId>`. PersistedChatSettings deliberately
@ -62,6 +66,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"
@ -262,9 +272,21 @@ 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;
codeToolsEnabled: boolean;
imageToolsEnabled: boolean;
/**
* Fetch pill state, independent of `toolsEnabled` (Search). Only
* consulted when `providerSupportsBuiltinWebFetch` is true.
*/
webFetchToolsEnabled: boolean;
toolStatus: string | null;
generatingStatus: string | null;
autoHealToolCalls: boolean;
@ -286,11 +308,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;
@ -324,6 +349,7 @@ type ChatRuntimeStore = {
setToolsEnabled: (enabled: boolean, options?: { persist?: boolean }) => void;
setCodeToolsEnabled: (enabled: boolean) => void;
setImageToolsEnabled: (enabled: boolean) => void;
setWebFetchToolsEnabled: (enabled: boolean) => void;
setToolStatus: (status: string | null) => void;
setGeneratingStatus: (status: string | null) => void;
setAutoHealToolCalls: (enabled: boolean) => void;
@ -336,6 +362,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;
};
@ -377,6 +407,7 @@ const PERSISTED_INFERENCE_PARAM_KEYS = [
"maxTokens",
"systemPrompt",
"trustRemoteCode",
"fastMode",
] as const satisfies readonly PersistedInferenceParamKey[];
const SCALAR_SETTING_KEYS = [
@ -564,9 +595,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
supportsBuiltinWebSearch: false,
supportsBuiltinCodeExecution: false,
supportsBuiltinImageGeneration: false,
supportsBuiltinWebFetch: false,
toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false),
codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false),
toolStatus: null,
generatingStatus: null,
autoHealToolCalls: true,
@ -587,6 +620,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
settingsPanelOpen: false,
pendingAudioBase64: null,
pendingAudioName: null,
pendingImageEditReference: null,
contextUsage: null,
modelLoading: false,
activeNativePathToken: null,
@ -640,7 +674,14 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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(() => {
@ -704,12 +745,37 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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) =>
@ -744,9 +810,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
supportsBuiltinWebSearch: false,
supportsBuiltinCodeExecution: false,
supportsBuiltinImageGeneration: false,
supportsBuiltinWebFetch: false,
toolsEnabled: false,
codeToolsEnabled: false,
imageToolsEnabled: false,
webFetchToolsEnabled: false,
toolStatus: null,
kvCacheDtype: null,
loadedKvCacheDtype: null,
@ -759,6 +827,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
defaultChatTemplate: null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
pendingImageEditReference: null,
}));
},
setReasoningEnabled: (reasoningEnabled, options) =>
@ -806,6 +875,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
saveBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, imageToolsEnabled);
return { imageToolsEnabled };
}),
setWebFetchToolsEnabled: (webFetchToolsEnabled) =>
set(() => {
saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled);
return { webFetchToolsEnabled };
}),
setToolStatus: (toolStatus) => set({ toolStatus }),
setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
setAutoHealToolCalls: (autoHealToolCalls) =>
@ -845,5 +919,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((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 }),
}));

40
studio/frontend/src/features/chat/types/api.ts
View file

@ -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,12 +193,31 @@ 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[];
export interface OpenAIChatMessage {
role: "system" | "user" | "assistant";
@ -262,6 +282,12 @@ 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 {

7
studio/frontend/src/features/chat/types/runtime.ts
View file

@ -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 {

5
studio/frontend/src/features/chat/utils/chat-settings-storage.ts
View file

@ -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;
}

80
studio/frontend/src/features/recipe-studio/components/inline/inline-model.tsx
View file

@ -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 })
}
/>
</InlineField>
<InlineField label="API key">
@ -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<string>();
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 (
<div className="grid gap-3 sm:grid-cols-2">
@ -82,12 +94,38 @@ export function InlineModel(props: InlineModelProps): ReactElement {
/>
</InlineField>
<InlineField label="Model">
<Input
className="nodrag h-8 w-full text-xs"
placeholder={isLinkedToLocal ? "local" : "gpt-4o-mini"}
value={modelConfig.model}
onChange={(event) => props.onUpdate({ model: event.target.value })}
/>
{isLinkedToLocal ? (
<LocalRecipeModelSelector
compact={true}
className="h-8 rounded-md text-xs"
value={
modelConfig.model.trim().toLowerCase() === "local"
? ""
: modelConfig.model
}
ggufVariant={modelConfig.gguf_variant}
onChange={(model, variant) =>
props.onUpdate({
model,
// biome-ignore lint/style/useNamingConvention: api schema
gguf_variant: variant ?? undefined,
})
}
/>
) : (
<Input
className="nodrag h-8 w-full text-xs"
placeholder="gpt-4o-mini"
value={modelConfig.model}
onChange={(event) =>
props.onUpdate({
model: event.target.value,
// biome-ignore lint/style/useNamingConvention: api schema
gguf_variant: undefined,
})
}
/>
)}
</InlineField>
<InlineField label="Temperature" className="sm:col-span-2">
<Input

644
studio/frontend/src/features/recipe-studio/dialogs/models/local-recipe-model-selector.tsx Normal file
View file

@ -0,0 +1,644 @@
// 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 { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import {
Popover,
PopoverContent,
PopoverTrigger,
} from "@/components/ui/popover";
import { Spinner } from "@/components/ui/spinner";
import {
type GgufVariantDetail,
type LocalModelInfo,
listGgufVariants,
listLocalModels,
} from "@/features/chat";
import { cn } from "@/lib/utils";
import { Link } from "@tanstack/react-router";
import { ChevronDownIcon, ChevronRightIcon, RefreshCwIcon } from "lucide-react";
import {
type ComponentPropsWithoutRef,
type ReactElement,
forwardRef,
useCallback,
useEffect,
useMemo,
useState,
} from "react";
const GGUF_SUFFIX_PATTERN = /-GGUF(?:$|-)/i;
type LocalRecipeModelSelectorProps = {
value: string;
ggufVariant?: string | null;
onChange: (modelId: string, ggufVariant?: string | null) => 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<GgufVariantDetail[] | null>(null);
const [defaultVariant, setDefaultVariant] = useState<string | null>(null);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(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 (
<div className="flex items-center gap-2 px-4 py-2 text-xs text-muted-foreground">
<Spinner className="size-3" />
Loading quantizations...
</div>
);
}
if (error) {
return <div className="px-4 py-2 text-xs text-destructive">{error}</div>;
}
if (!sortedVariants || sortedVariants.length === 0) {
return (
<div className="px-4 py-2 text-xs text-muted-foreground">
No GGUF quantizations found for this model.
</div>
);
}
return (
<div className="ml-6 mt-1 rounded-lg bg-muted/25 p-1.5">
<div className="mb-1 px-2 text-[10px] font-medium uppercase tracking-wide text-muted-foreground">
Quantization
</div>
<div className="space-y-0.5">
{sortedVariants.map((variant) => {
const selected = selectedVariant === variant.quant;
return (
<button
key={variant.filename}
type="button"
onClick={() => onSelect(variant.quant)}
className={cn(
"flex w-full items-center gap-2 rounded-md px-2 py-1.5 text-left text-xs transition-colors hover:bg-muted/60",
selected && "bg-background text-foreground shadow-sm",
)}
>
<span className="min-w-0 flex-1 truncate font-mono">
{variant.quant}
</span>
{variant.quant === defaultVariant ? (
<Badge variant="secondary" className="h-4 px-1.5 text-[10px]">
recommended
</Badge>
) : null}
{variant.downloaded ? (
<Badge variant="outline" className="h-4 px-1.5 text-[10px]">
ready
</Badge>
) : null}
</button>
);
})}
</div>
</div>
);
}
type SelectorTriggerProps = ComponentPropsWithoutRef<"button"> & {
value: string;
selectedModel: LocalModelInfo | null;
ggufVariant?: string | null;
inputId?: string;
disabled: boolean;
compact: boolean;
className?: string;
};
const SelectorTrigger = forwardRef<HTMLButtonElement, SelectorTriggerProps>(
function SelectorTrigger(
{
value,
selectedModel,
ggufVariant,
inputId,
disabled,
compact,
className,
...triggerProps
},
ref,
): ReactElement {
const selected = getSelectedModelSummary(value, selectedModel, ggufVariant);
return (
<button
{...triggerProps}
ref={ref}
id={inputId}
type="button"
disabled={disabled}
className={cn(
"nodrag flex w-full min-w-0 items-center gap-2 rounded-xl border border-border/70 bg-background px-3 text-left transition-colors hover:bg-muted/40 disabled:pointer-events-none disabled:opacity-60",
compact ? "min-h-8 py-1.5 text-xs" : "min-h-10 py-2 text-sm",
className,
)}
>
<span className="min-w-0 flex-1">
<span
className={cn(
"block truncate font-medium",
!selected.label && "text-muted-foreground",
)}
>
{selected.label || "Choose a local model"}
</span>
{compact ? null : (
<span className="mt-0.5 flex min-w-0 items-center gap-1.5 text-[11px] text-muted-foreground">
<span className="truncate">
{selected.label
? selected.source
: "Select from local and cached models"}
</span>
{selected.isGguf ? <span>GGUF</span> : null}
{ggufVariant ? (
<span className="truncate font-mono">{ggufVariant}</span>
) : null}
</span>
)}
</span>
{compact && ggufVariant ? (
<Badge
variant="secondary"
className="h-4 px-1.5 font-mono text-[10px]"
>
{ggufVariant}
</Badge>
) : null}
<ChevronDownIcon className="size-4 shrink-0 text-muted-foreground" />
</button>
);
},
);
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 (
<div>
<button
type="button"
disabled={probing}
onClick={() => onSelectModel(model)}
className={cn(
"flex w-full items-center gap-2 rounded-lg px-2.5 py-2.5 text-left text-sm transition-colors hover:bg-muted/50",
selected && "bg-muted/70 text-foreground ring-1 ring-border/70",
)}
>
{expandable ? (
expanded ? (
<ChevronDownIcon className="size-3.5 shrink-0 text-muted-foreground" />
) : (
<ChevronRightIcon className="size-3.5 shrink-0 text-muted-foreground" />
)
) : (
<span className="size-3.5 shrink-0" />
)}
<span className="min-w-0 flex-1">
<span className="block truncate font-medium">
{getModelLabel(model)}
</span>
<span className="mt-0.5 block truncate text-[11px] text-muted-foreground">
{model.id}
</span>
</span>
<span className="flex shrink-0 items-center gap-1">
{probing ? (
<Spinner className="size-3 text-muted-foreground" />
) : null}
{expandable || directGguf ? (
<Badge variant="secondary" className="h-4 px-1.5 text-[10px]">
GGUF
</Badge>
) : null}
<Badge variant="outline" className="h-4 px-1.5 text-[10px]">
{sourceLabel(model)}
</Badge>
</span>
</button>
{expanded ? (
<LocalGgufVariantList
repoId={model.id}
selectedVariant={selected ? ggufVariant : null}
onSelect={(variant) => onSelectVariant(model.id, variant)}
/>
) : null}
</div>
);
}
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 (
<div className="flex items-center gap-2 px-3 py-3 text-xs text-muted-foreground">
<Spinner className="size-3" />
Scanning local models...
</div>
);
}
if (error) {
return (
<div className="space-y-2 px-3 py-3 text-xs">
<p className="text-destructive">{error}</p>
<Button type="button" variant="outline" size="xs" onClick={onRefresh}>
Try again
</Button>
</div>
);
}
if (models.length === 0) {
return (
<div className="space-y-2 px-3 py-3 text-xs text-muted-foreground">
<p className="font-medium text-foreground">No local models found.</p>
<p>
Download a model or add a scan folder from Chat, then refresh this
list.
</p>
<Link
to="/chat"
className="inline-flex font-medium text-primary underline-offset-4 hover:underline"
>
Open Chat model picker
</Link>
</div>
);
}
return (
<div className="space-y-1">
{models.map((model) => (
<LocalModelRow
key={model.id}
model={model}
selected={model.id === value}
expanded={expandedModelId === model.id}
probing={probingVariantModelId === model.id}
ggufVariant={ggufVariant}
onSelectModel={onSelectModel}
onSelectVariant={onSelectVariant}
/>
))}
</div>
);
}
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<LocalModelInfo[]>([]);
const [loading, setLoading] = useState(false);
const [error, setError] = useState<string | null>(null);
const [expandedModelId, setExpandedModelId] = useState<string | null>(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 (
<Popover open={open} onOpenChange={handleOpenChange}>
<PopoverTrigger asChild={true}>
<SelectorTrigger
value={value}
selectedModel={selectedModel}
ggufVariant={ggufVariant}
inputId={inputId}
disabled={disabled}
compact={compact}
className={className}
/>
</PopoverTrigger>
<PopoverContent
align="start"
sideOffset={6}
className="menu-soft-surface nodrag nowheel gap-0 overflow-hidden p-0"
style={{
width:
"min(max(var(--radix-popover-trigger-width), 34rem), calc(100vw - 1rem))",
}}
>
<div className="flex flex-col">
<div className="border-b border-border/60 p-2.5">
<div className="flex items-center gap-2">
<Input
value={query}
onChange={(event) => setQuery(event.target.value)}
placeholder="Filter local models"
className="h-8 flex-1"
autoFocus={true}
/>
<Button
type="button"
variant="ghost"
size="icon-sm"
onClick={requestModelRefresh}
aria-label="Refresh local models"
>
<RefreshCwIcon className="size-3.5" />
</Button>
</div>
</div>
<div
className="nowheel max-h-[min(24rem,calc(100vh-12rem))] overflow-y-auto overscroll-contain p-1.5"
onWheelCapture={(event) => event.stopPropagation()}
>
<LocalModelResults
loading={loading}
error={error}
models={filteredModels}
value={value}
ggufVariant={ggufVariant}
expandedModelId={expandedModelId}
probingVariantModelId={probingVariantModelId}
onRefresh={requestModelRefresh}
onSelectModel={selectModel}
onSelectVariant={selectVariant}
/>
</div>
</div>
</PopoverContent>
</Popover>
);
}

86
studio/frontend/src/features/recipe-studio/dialogs/models/model-config-dialog.tsx
View file

@ -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<HTMLDivElement>(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<ModelConfig>);
};
// 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
</p>
<p className="mt-1 text-xs text-muted-foreground">
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.
</p>
</div>
<div className="grid gap-1.5">
@ -144,15 +155,48 @@ export function ModelConfigDialog({
<FieldLabel
label="Model ID"
htmlFor={modelId}
hint={isLinkedToLocal ? "Uses the model loaded in Chat. Any value works here." : "The exact model name sent to the connection."}
/>
<Input
id={modelId}
className="nodrag"
placeholder={isLinkedToLocal ? "local" : "gpt-4o-mini"}
value={config.model}
onChange={(event) => 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 ? (
<LocalRecipeModelSelector
inputId={modelId}
value={
config.model.trim().toLowerCase() === "local" ? "" : config.model
}
ggufVariant={config.gguf_variant}
onChange={(model, variant) =>
onUpdate({
model,
// biome-ignore lint/style/useNamingConvention: api schema
gguf_variant: variant ?? undefined,
})
}
/>
) : (
<Input
id={modelId}
className="nodrag"
placeholder="gpt-4o-mini"
value={config.model}
onChange={(event) =>
onUpdate({
model: event.target.value,
// biome-ignore lint/style/useNamingConvention: api schema
gguf_variant: undefined,
})
}
/>
)}
{isLinkedToLocal ? (
<p className="text-xs text-muted-foreground">
Recipes will load this model automatically. GGUF quantization is
saved with the preset.
</p>
) : null}
</div>
<div className="grid gap-3">
<div className="space-y-1">
@ -250,8 +294,12 @@ export function ModelConfigDialog({
}
/>
</div>
<label className="flex items-center gap-2 text-xs font-semibold uppercase text-muted-foreground">
<label
htmlFor={skipHealthCheckId}
className="flex items-center gap-2 text-xs font-semibold uppercase text-muted-foreground"
>
<Checkbox
id={skipHealthCheckId}
checked={config.skip_health_check ?? false}
onCheckedChange={(value) =>
updateField("skip_health_check", Boolean(value))

85
studio/frontend/src/features/recipe-studio/easy/github-crawler-easy-view.tsx
View file

@ -1,13 +1,14 @@
// 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 { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { GithubIcon, PlayCircleIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { type ReactElement, useEffect, useMemo, useState } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { FieldLabel } from "../dialogs/shared/field-label";
import { LocalRecipeModelSelector } from "../dialogs/models/local-recipe-model-selector";
import { GithubRepoSeedForm } from "../dialogs/seed/seed-dialog";
import { FieldLabel } from "../dialogs/shared/field-label";
import type { ModelConfig, NodeConfig, SeedConfig } from "../types";
type GithubCrawlerEasyViewProps = {
@ -44,6 +45,21 @@ export function GithubCrawlerEasyView({
) ?? null,
[configs],
);
const localProviderNames = useMemo(() => {
const names = new Set<string>();
for (const config of Object.values(configs)) {
if (config.kind === "model_provider" && config.is_local === true) {
names.add(config.name);
}
}
return names;
}, [configs]);
const isModelLinkedToLocal = modelConfig
? localProviderNames.has(modelConfig.provider)
: false;
const modelValue = modelConfig?.model ?? "";
const localModelValue =
modelValue.trim().toLowerCase() === "local" ? "" : modelValue;
// Local buffer for the Rows input so the user can hold transient invalid
// state (empty while backspacing, partial digits, etc.) without the parent
@ -52,17 +68,40 @@ export function GithubCrawlerEasyView({
// blur we clamp back to a sane default if the user left it empty.
const [rowsText, setRowsText] = useState(String(rows));
useEffect(() => {
// eslint-disable-next-line react-hooks/set-state-in-effect -- keep the draft input in sync when the parent resets rows.
setRowsText(String(rows));
}, [rows]);
const handleSeedUpdate = (patch: Partial<SeedConfig>): void => {
if (!seedConfig) return;
if (!seedConfig) {
return;
}
updateConfig(seedConfig.id, patch);
};
const handleModelChange = (value: string): void => {
if (!modelConfig) return;
updateConfig(modelConfig.id, { model: value });
if (!modelConfig) {
return;
}
updateConfig(modelConfig.id, {
model: value,
// biome-ignore lint/style/useNamingConvention: api schema
gguf_variant: undefined,
});
};
const handleLocalModelChange = (
model: string,
variant?: string | null,
): void => {
if (!modelConfig) {
return;
}
updateConfig(modelConfig.id, {
model,
// biome-ignore lint/style/useNamingConvention: api schema
gguf_variant: variant ?? undefined,
});
};
if (!seedConfig) {
@ -96,9 +135,8 @@ export function GithubCrawlerEasyView({
<h2 className="text-base font-semibold">GitHub Crawler</h2>
<p className="text-xs text-muted-foreground">
Crawl real GitHub issues and PRs and turn each thread into a{" "}
<code>{"{User, Assistant}"}</code> training pair.
Defaults use the server's <code>GH_TOKEN</code> env var and the
bundled local model.
<code>{"{User, Assistant}"}</code> training pair. Defaults use the
server's <code>GH_TOKEN</code> env var and the bundled local model.
</p>
</div>
</div>
@ -149,15 +187,28 @@ export function GithubCrawlerEasyView({
<div className="grid gap-1.5">
<FieldLabel
label="Model"
hint="OpenAI-compatible model id. Local GGUFs run on the bundled llama-server."
/>
<Input
className="nodrag font-mono text-xs"
value={modelConfig?.model ?? ""}
onChange={(event) => 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 ? (
<LocalRecipeModelSelector
value={localModelValue}
ggufVariant={modelConfig?.gguf_variant}
onChange={handleLocalModelChange}
disabled={!modelConfig}
/>
) : (
<Input
className="nodrag font-mono text-xs"
value={modelValue}
onChange={(event) => handleModelChange(event.target.value)}
placeholder="unsloth/gemma-4-E2B-it-GGUF"
disabled={!modelConfig}
/>
)}
</div>
</div>
</section>

101
studio/frontend/src/features/recipe-studio/executions/tracker.ts
View file

@ -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<boolean> {
}: TrackRecipeExecutionParams): Promise<TrackRecipeExecutionResult> {
let done = false;
let lastStatus: RecipeExecutionStatus = initialExecution.status;
let completedEventPayload: Record<string, unknown> | 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<string, unknown>)
completedPayload &&
typeof completedPayload.processor_artifacts === "object" &&
completedPayload.processor_artifacts !== null
? (completedPayload.processor_artifacts as Record<string, unknown>)
: 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 };
}

565
studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts
View file

@ -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<string | null> {
const GGUF_MODEL_PATTERN = /gguf/i;
function collectUsedLlmModelAliases(payload: RecipePayload): Set<string> {
const columns = Array.isArray(payload.recipe.columns)
? payload.recipe.columns
: [];
const aliases = new Set<string>();
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<string> {
const providers = Array.isArray(payload.recipe.model_providers)
? (payload.recipe.model_providers as Array<Record<string, unknown>>)
? (payload.recipe.model_providers as Record<string, unknown>[])
: [];
const localProviderNames = new Set<string>();
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<string>,
): Record<string, unknown>[] {
const usedAliases = collectUsedLlmModelAliases(payload);
if (usedAliases.size === 0) {
return [];
}
const modelConfigs = Array.isArray(payload.recipe.model_configs)
? (payload.recipe.model_configs as Array<Record<string, unknown>>)
? 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<string, unknown>,
): 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<string, unknown>[],
): LocalModelLoadPlan | null {
const selections = new Map<string, LocalModelSelection>();
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<boolean> {
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<string | null> {
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<LocalModelSelection | null> {
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<RestorableLocalModelSnapshot> {
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<string | null> {
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<void>) | 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<void>) | null;
}): Promise<boolean> => {
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<void>) | 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<typeof validateRecipe>[0],
): Promise<boolean> => {
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<boolean> => {
const runPreview = useCallback((): Promise<boolean> => {
return runWithValidation("preview", previewRows, null);
}, [previewRows, runWithValidation]);
const runFull = useCallback(async (): Promise<boolean> => {
const runFull = useCallback((): Promise<boolean> => {
return runWithValidation("full", fullRows, fullRunName);
}, [fullRows, fullRunName, runWithValidation]);
const runFromDialog = useCallback(async (): Promise<boolean> => {
const runFromDialog = useCallback((): Promise<boolean> => {
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<void> => {
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,

20
studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts
View file

@ -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),
);
}

98
studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts
View file

@ -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<string, NodeConfig>): boolean {
function isModelSemanticEdge(
edge: Edge,
configs: Record<string, NodeConfig>,
): boolean {
const source = configs[edge.source];
const target = configs[edge.target];
return Boolean(
@ -315,12 +322,16 @@ export const useRecipeStudioStore = create<RecipeStudioState>((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<RecipeStudioState>((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<RecipeStudioState>((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<RecipeStudioState>((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<RecipeStudioState>((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<RecipeStudioState>((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<RecipeStudioState>((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,
},
};
}
}
}

2
studio/frontend/src/features/recipe-studio/types/index.ts
View file

@ -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;

65
studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts
View file

@ -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 } };
}

13
studio/frontend/src/features/recipe-studio/utils/import/parsers/model-parser.ts
View file

@ -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<string, unknown>,
@ -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),

73
studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts
View file

@ -54,55 +54,62 @@ export function buildModelProvider(
};
}
export function buildModelConfig(
function assignFiniteNumber(
target: Record<string, unknown>,
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<string, unknown> {
const inference: Record<string, unknown> = {};
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<string, unknown> {
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:

31
studio/frontend/src/features/recipe-studio/utils/payload/validate.ts
View file

@ -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 (

60
studio/frontend/src/index.css
View file

@ -1188,6 +1188,53 @@
border-color: var(--border) !important;
}
.generated-image-loading-card {
position: relative;
overflow: hidden;
contain: paint;
}
.generated-image-loading-wave {
position: relative;
display: grid;
grid-template-columns: repeat(8, minmax(0, 1fr));
gap: 14px;
width: min(66%, 18rem);
padding: 1.5rem;
border-radius: 1.5rem;
}
.generated-image-loading-dot {
width: 7px;
height: 7px;
border-radius: 9999px;
background: color-mix(in oklch, var(--muted-foreground) 82%, var(--primary));
opacity: 0.12;
transform: translate3d(0, 4px, 0) scale(0.72);
animation: generated-image-dot-wave 1850ms var(--ease-out-quart) infinite;
animation-delay: calc((var(--dot-row) * 72ms) + (var(--dot-col) * 72ms));
will-change: transform, opacity;
}
@keyframes generated-image-dot-wave {
0%,
22%,
100% {
opacity: 0.1;
transform: translate3d(0, 4px, 0) scale(0.72);
}
46% {
opacity: 0.46;
transform: translate3d(0, -3px, 0) scale(0.96);
}
66% {
opacity: 0.2;
transform: translate3d(0, 0, 0) scale(0.82);
}
}
/*
* prefers-reduced-motion: honour the OS-level "reduce motion" preference.
* Tailwind animate-in/out, Radix open/close transforms, infinite shine/pulse
@ -1197,11 +1244,11 @@
* end state. Hover colour changes become instant rather than fading, which is
* the documented WCAG outcome (motion is "minimised, not removed").
*
* .animate-spin is the exception: loading spinners are essential progress
* .animate-spin and generated image loading dots are the exceptions: loading
* indicators across Studio (tool execution loaders, sonner toasts, Tauri
* startup / update screens, the <Spinner /> primitive). Freezing them
* removes the only visual signal that work is in flight, so they keep
* animating but at a slower, less aggressive 1.5s cadence.
* startup / update screens, the <Spinner /> primitive, and image generation
* cards). Freezing them removes the only visual signal that work is in flight,
* so they keep animating.
*/
@media (prefers-reduced-motion: reduce) {
*,
@ -1217,4 +1264,9 @@
animation-duration: 1.5s !important;
animation-iteration-count: infinite !important;
}
.generated-image-loading-dot {
animation-duration: 1850ms !important;
animation-iteration-count: infinite !important;
}
}

21
studio/install_python_stack.py
View file

@ -37,6 +37,7 @@ from backend.utils.wheel_utils import (
IS_WINDOWS = sys.platform == "win32"
IS_MACOS = sys.platform == "darwin"
IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64"
IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64"
# ── ROCm / AMD GPU support ─────────────────────────────────────────────────────
# Mapping from detected ROCm (major, minor) to the best PyTorch wheel tag on
@ -423,6 +424,7 @@ def _infer_no_torch() -> bool:
NO_TORCH = _infer_no_torch()
# -- Verbosity control ----------------------------------------------------------
# By default the installer shows a minimal progress bar (one line, in-place).
# Set UNSLOTH_VERBOSE=1 in the environment to restore full per-step output:
@ -448,6 +450,11 @@ LOCAL_DD_GITHUB_PLUGIN = (
SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed"
)
# Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides file).
_MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt"
if IS_MAC_ARM and _MLX_OVERRIDES.is_file():
os.environ.setdefault("UV_OVERRIDE", str(_MLX_OVERRIDES))
# -- Unicode-safe printing ---------------------------------------------
# On Windows the default console encoding can be a legacy code page
# (e.g. CP1252) that cannot represent Unicode glyphs such as ✅ or ❌.
@ -960,6 +967,20 @@ def install_python_stack() -> int:
[sys.executable, "-m", "pip", "install", "--upgrade", "pip"],
)
# macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the
# mlx-vlm / mlx-lm transformers pin -- set at module load).
if IS_MAC_ARM and not skip_base:
_progress("MLX stack (Apple Silicon)")
pip_install(
"Installing MLX stack (mlx + mlx-lm + mlx-vlm)",
"--no-cache-dir",
"--upgrade",
"mlx",
"mlx-metal",
"mlx-lm",
"mlx-vlm",
)
# 3. Core packages: unsloth-zoo + unsloth (or custom package name)
if skip_base:
pass

77
tests/python/test_construct_chat_template_validation.py Normal file
View file

@ -0,0 +1,77 @@
"""Negative-path validation tests for unsloth.chat_templates.construct_chat_template.
Regression coverage for the str.find() / regex no-match guards added in
PR #5763 follow-up: missing placeholders or unrecoverable two-example
structures must raise RuntimeError with a clear message, not IndexError
or AttributeError, and must never silently drop the last character via
s[:-1].
Uses a minimal fake tokenizer so the cases run on CPU-only CI without
HF_TOKEN and without downloading a gated model. The validation paths
exercised here fail before construct_chat_template reaches any heavy
tokenizer interaction, so the stub stays small.
"""
import pytest
from unsloth.chat_templates import construct_chat_template
class _FakeTokenizer:
"""Minimum surface construct_chat_template touches before the
validation guards fire."""
name_or_path = "fake/tokenizer"
eos_token = "</s>"
def get_vocab(self):
return {"</s>": 0}
@pytest.mark.parametrize(
"template, expected_in_message",
[
("only {INPUT} here, no output marker", "{OUTPUT}"),
("only {OUTPUT} here, no input marker", "{INPUT}"),
("neither sentinel here, just literal text", "{INPUT}"),
("neither sentinel here, just literal text", "{OUTPUT}"),
],
)
def test_missing_placeholder_in_chat_template_raises(template, expected_in_message):
with pytest.raises(RuntimeError) as exc_info:
construct_chat_template(
tokenizer = _FakeTokenizer(),
chat_template = template,
extra_eos_tokens = ["</s>"],
)
assert expected_in_message in str(exc_info.value)
def test_single_pair_template_raises_clear_error_not_attribute_error():
"""One {INPUT}/{OUTPUT} pair (rather than the required two) used to
crash with AttributeError on `found.group(1)` after the for-loop
broke without setting `found`. Must raise RuntimeError now."""
template = "user: {INPUT}\nassistant: {OUTPUT}\n"
with pytest.raises(RuntimeError):
construct_chat_template(
tokenizer = _FakeTokenizer(),
chat_template = template,
extra_eos_tokens = ["</s>"],
)
def test_error_message_excerpt_is_bounded():
"""Error messages must include a bounded excerpt of the offending
template, not dump arbitrarily large content into the traceback."""
huge = ("garbage " * 5000) + "{INPUT}" # ~40 KB, missing {OUTPUT}
with pytest.raises(RuntimeError) as exc_info:
construct_chat_template(
tokenizer = _FakeTokenizer(),
chat_template = huge,
extra_eos_tokens = ["</s>"],
)
msg = str(exc_info.value)
# Excerpt is repr-quoted and capped; total message should stay well
# under the template length.
assert len(msg) < 1000
assert "{OUTPUT}" in msg

41
unsloth/chat_templates.py
View file

@ -2461,17 +2461,40 @@ extra_eos_tokens = None,
f"{left_changed}"
)
except:
ending = chat_template[chat_template.find("{OUTPUT}") + len("{OUTPUT}"):]
output_pos = chat_template.find("{OUTPUT}")
input_pos = chat_template.find("{INPUT}")
if output_pos == -1 or input_pos == -1:
missing = []
if input_pos == -1: missing.append("{INPUT}")
if output_pos == -1: missing.append("{OUTPUT}")
raise RuntimeError(
f"Unsloth: chat_template must contain {' and '.join(missing)} "
f"placeholder(s). Got: {chat_template[:200]!r}"
)
ending = chat_template[output_pos + len("{OUTPUT}"):]
ending = re.escape(ending)
find_text = "{INPUT}" + ending + "(.+?{OUTPUT}" + ending + ")"
response_part = re.findall(find_text, chat_template, flags = re.DOTALL | re.MULTILINE)
if len(response_part) == 0:
raise RuntimeError(
"Unsloth: Could not recover a two-example structure from chat_template. "
"Provide exactly two {INPUT}/{OUTPUT} pairs (and optionally {SYSTEM}). "
f"Got: {chat_template[:200]!r}"
)
response_part = response_part[0]
found = None
for j in range(1, len(response_part)):
try_find = re.escape(response_part[:j])
try: found = next(re.finditer("(" + try_find + ").+?\\{INPUT\\}", chat_template, flags = re.DOTALL | re.MULTILINE))
except: break
if found is None:
raise RuntimeError(
"Unsloth: Could not locate a separator between examples in chat_template. "
"Provide exactly two {INPUT}/{OUTPUT} pairs (and optionally {SYSTEM}). "
f"Got: {chat_template[:200]!r}"
)
separator = found.group(1)
response_start = chat_template.find(response_part)
@ -2607,8 +2630,20 @@ extra_eos_tokens = None,
jinja_template = "{{ bos_token }}" + jinja_template
# Get instruction and output parts for train_on_inputs = False
input_part = input_part [:input_part .find("{INPUT}")]
output_part = output_part[:output_part.find("{OUTPUT}")]
input_idx = input_part .find("{INPUT}")
output_idx = output_part.find("{OUTPUT}")
if input_idx == -1:
raise RuntimeError(
f"Unsloth: The instruction section of the template must contain the "
f"'{{INPUT}}' placeholder. Section: {input_part[:200]!r}"
)
if output_idx == -1:
raise RuntimeError(
f"Unsloth: The response section of the template must contain the "
f"'{{OUTPUT}}' placeholder. Section: {output_part[:200]!r}"
)
input_part = input_part [:input_idx ]
output_part = output_part[:output_idx]
return modelfile, jinja_template, input_part, output_part

2
unsloth/models/_utils.py
View file

@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2026.5.7"
__version__ = "2026.5.8"
__all__ = [
"SUPPORTS_BFLOAT16",

83
unsloth_cli/commands/studio.py
View file

@ -206,6 +206,79 @@ def _find_setup_script() -> Optional[Path]:
return None
def _iter_editable_studio_source_roots(venv_dir: Path):
"""Yield repo roots from setuptools `__editable___*_finder.py` files in
*venv_dir*'s site-packages whose MAPPING includes a `studio` entry.
Returns the parent dir of the mapped `studio` package (i.e. the repo
root), so callers can append `/studio/...` to reach any subdir.
"""
import ast
import re
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:
yield Path(studio_pkg).parent
def _find_frontend_dist() -> Optional[Path]:
"""Locate a built `studio/frontend/dist` (containing index.html).
Probes (in order): package-local default, installer venv site-packages,
editable source roots referenced from the installer venv. Returns None
if nothing servable is found, so callers can decide to error or proceed
in `--api-only` mode.
Fixes the silent 404 when another `unsloth` on PATH shadows the
installer's binary and points `_PACKAGE_ROOT` at a site-packages copy
that never received a vite build.
"""
candidates: List[Path] = [_PACKAGE_ROOT / "studio" / "frontend" / "dist"]
venv_dir = STUDIO_HOME / "unsloth_studio"
for pattern in (
"lib/python*/site-packages/studio/frontend/dist",
"Lib/site-packages/studio/frontend/dist",
):
candidates.extend(venv_dir.glob(pattern))
for repo_root in _iter_editable_studio_source_roots(venv_dir):
candidates.append(repo_root / "studio" / "frontend" / "dist")
seen: set[Path] = set()
for c in candidates:
try:
resolved = c.resolve()
except OSError:
resolved = c
if resolved in seen:
continue
seen.add(resolved)
if (c / "index.html").is_file():
return c
return None
# ── helpers for `unsloth studio run` ────────────────────────────────
@ -539,8 +612,14 @@ def studio_default(
"--port",
str(port),
]
if frontend:
args.extend(["--frontend", str(frontend)])
# Resolve frontend explicitly so the spawned run.py uses a real
# built dist regardless of where its __file__ lands. Skip in
# --api-only (no UI served).
resolved_frontend = frontend
if resolved_frontend is None and not api_only:
resolved_frontend = _find_frontend_dist()
if resolved_frontend is not None:
args.extend(["--frontend", str(resolved_frontend)])
if silent:
args.append("--silent")
if api_only: